Honor stored tax-benefit model versions at runtime (#220)

* Honor stored model versions at runtime

* Validate shared runtime model bundles

* Harden shared runtime bundle resolution

Author: MaxGhenis

changelog.d/runtime-version-resolution.changed.md

@@ -0,0 +1 @@
+Reject shared runtime bundle reuse when the compared database rows point at different model identities, even if their runtime version strings match.

src/policyengine_api/api/analysis.py

@@ -64,6 +64,9 @@
     TaxBenefitModel,
     TaxBenefitModelVersion,
 )
+from policyengine_api.runtime_versions import (
+    resolve_shared_runtime_model_version_from_db,
+)
 from policyengine_api.services.database import get_session
 from policyengine_api.services.model_resolver import (
     resolve_country_from_simulation,
@@ -752,7 +755,6 @@ def _run_local_economy_comparison_uk(
     from policyengine.core.dynamic import Dynamic as PEDynamic
     from policyengine.core.policy import ParameterValue as PEParameterValue
     from policyengine.core.policy import Policy as PEPolicy
-    from policyengine.tax_benefit_models.uk import uk_latest
     from policyengine.tax_benefit_models.uk.datasets import PolicyEngineUKDataset
 
     from policyengine_api.models import Policy as DBPolicy
@@ -778,7 +780,11 @@ def _run_local_economy_comparison_uk(
     if not dataset:
         raise ValueError(f"Dataset {baseline_sim.dataset_id} not found")
 
-    pe_model_version = uk_latest
+    pe_model_version = resolve_shared_runtime_model_version_from_db(
+        session,
+        baseline_sim.tax_benefit_model_version_id,
+        reform_sim.tax_benefit_model_version_id,
+    )
     param_lookup = {p.name: p for p in pe_model_version.parameters}
 
     def build_policy(policy_id):
@@ -937,7 +943,6 @@ def _run_local_economy_comparison_us(
     from policyengine.core.dynamic import Dynamic as PEDynamic
     from policyengine.core.policy import ParameterValue as PEParameterValue
     from policyengine.core.policy import Policy as PEPolicy
-    from policyengine.tax_benefit_models.us import us_latest
     from policyengine.tax_benefit_models.us.datasets import PolicyEngineUSDataset
 
     from policyengine_api.models import Policy as DBPolicy
@@ -963,7 +968,11 @@ def _run_local_economy_comparison_us(
     if not dataset:
         raise ValueError(f"Dataset {baseline_sim.dataset_id} not found")
 
-    pe_model_version = us_latest
+    pe_model_version = resolve_shared_runtime_model_version_from_db(
+        session,
+        baseline_sim.tax_benefit_model_version_id,
+        reform_sim.tax_benefit_model_version_id,
+    )
     param_lookup = {p.name: p for p in pe_model_version.parameters}
 
     def build_policy(policy_id):

src/policyengine_api/modal/functions/__init__.py

@@ -15,6 +15,10 @@
     get_database_url,
     get_db_session,
 )
+from policyengine_api.runtime_versions import (
+    resolve_runtime_model_version_from_db,
+    resolve_shared_runtime_model_version_from_db,
+)
 
 # Required environment variables from each secret
 REQUIRED_DB_VARS = ["DATABASE_URL", "SUPABASE_URL", "SUPABASE_KEY"]
@@ -674,12 +678,13 @@ def simulate_economy_uk(simulation_id: str, traceparent: str | None = None) -> N
 
                     # Import policyengine
                     from policyengine.core import Simulation as PESimulation
-                    from policyengine.tax_benefit_models.uk import uk_latest
                     from policyengine.tax_benefit_models.uk.datasets import (
                         PolicyEngineUKDataset,
                     )
 
-                    pe_model_version = uk_latest
+                    pe_model_version = resolve_runtime_model_version_from_db(
+                        session, simulation.tax_benefit_model_version_id
+                    )
 
                     # Get policy and dynamic
                     policy = _get_pe_policy_uk(
@@ -847,12 +852,13 @@ def simulate_economy_us(simulation_id: str, traceparent: str | None = None) -> N
 
                     # Import policyengine
                     from policyengine.core import Simulation as PESimulation
-                    from policyengine.tax_benefit_models.us import us_latest
                     from policyengine.tax_benefit_models.us.datasets import (
                         PolicyEngineUSDataset,
                     )
 
-                    pe_model_version = us_latest
+                    pe_model_version = resolve_runtime_model_version_from_db(
+                        session, simulation.tax_benefit_model_version_id
+                    )
 
                     # Get policy and dynamic
                     policy = _get_pe_policy_us(
@@ -1008,7 +1014,6 @@ def economy_comparison_uk(job_id: str, traceparent: str | None = None) -> None:
                     ReportStatus,
                     Simulation,
                     SimulationStatus,
-                    TaxBenefitModelVersion,
                 )
 
                 with Session(engine) as session:
@@ -1035,12 +1040,6 @@ def economy_comparison_uk(job_id: str, traceparent: str | None = None) -> None:
                     if not dataset:
                         raise ValueError(f"Dataset {baseline_sim.dataset_id} not found")
 
-                    # Get model version (unused but keeping for reference)
-                    _ = session.get(
-                        TaxBenefitModelVersion,
-                        baseline_sim.tax_benefit_model_version_id,
-                    )
-
                     # Import policyengine
                     from policyengine.core import Simulation as PESimulation
                     from policyengine.outputs import DecileImpact as PEDecileImpact
@@ -1050,15 +1049,18 @@ def economy_comparison_uk(job_id: str, traceparent: str | None = None) -> None:
                     from policyengine.outputs.aggregate import (
                         AggregateType as PEAggregateType,
                     )
-                    from policyengine.tax_benefit_models.uk import uk_latest
                     from policyengine.tax_benefit_models.uk.datasets import (
                         PolicyEngineUKDataset,
                     )
                     from policyengine.tax_benefit_models.uk.outputs import (
                         ProgrammeStatistics as PEProgrammeStats,
                     )
 
-                    pe_model_version = uk_latest
+                    pe_model_version = resolve_shared_runtime_model_version_from_db(
+                        session,
+                        baseline_sim.tax_benefit_model_version_id,
+                        reform_sim.tax_benefit_model_version_id,
+                    )
 
                     # Get policies
                     baseline_policy = _get_pe_policy_uk(
@@ -1725,15 +1727,18 @@ def economy_comparison_us(job_id: str, traceparent: str | None = None) -> None:
                     from policyengine.outputs.aggregate import (
                         AggregateType as PEAggregateType,
                     )
-                    from policyengine.tax_benefit_models.us import us_latest
                     from policyengine.tax_benefit_models.us.datasets import (
                         PolicyEngineUSDataset,
                     )
                     from policyengine.tax_benefit_models.us.outputs import (
                         ProgramStatistics as PEProgramStats,
                     )
 
-                    pe_model_version = us_latest
+                    pe_model_version = resolve_shared_runtime_model_version_from_db(
+                        session,
+                        baseline_sim.tax_benefit_model_version_id,
+                        reform_sim.tax_benefit_model_version_id,
+                    )
 
                     # Get policies
                     baseline_policy = _get_pe_policy_us(
@@ -2631,7 +2636,6 @@ def compute_aggregate_uk(aggregate_id: str, traceparent: str | None = None) -> N
                 from policyengine.core import Simulation as PESimulation
                 from policyengine.outputs import Aggregate as PEAggregate
                 from policyengine.outputs import AggregateType as PEAggregateType
-                from policyengine.tax_benefit_models.uk import uk_latest
                 from policyengine.tax_benefit_models.uk.datasets import (
                     PolicyEngineUKDataset,
                 )
@@ -2692,6 +2696,10 @@ def compute_aggregate_uk(aggregate_id: str, traceparent: str | None = None) -> N
                         )
 
                     # Create policyengine simulation with loaded output
+                    pe_model_version = resolve_runtime_model_version_from_db(
+                        session, simulation.tax_benefit_model_version_id
+                    )
+
                     with logfire.span("load_output"):
                         pe_output_dataset = PolicyEngineUKDataset(
                             name=output_dataset.name or "output",
@@ -2703,7 +2711,7 @@ def compute_aggregate_uk(aggregate_id: str, traceparent: str | None = None) -> N
 
                         pe_sim = PESimulation(
                             dataset=pe_output_dataset,  # Use output as dataset
-                            tax_benefit_model_version=uk_latest,
+                            tax_benefit_model_version=pe_model_version,
                         )
                         pe_sim.output_dataset = pe_output_dataset
 
@@ -2787,7 +2795,6 @@ def compute_aggregate_us(aggregate_id: str, traceparent: str | None = None) -> N
                 from policyengine.core import Simulation as PESimulation
                 from policyengine.outputs import Aggregate as PEAggregate
                 from policyengine.outputs import AggregateType as PEAggregateType
-                from policyengine.tax_benefit_models.us import us_latest
                 from policyengine.tax_benefit_models.us.datasets import (
                     PolicyEngineUSDataset,
                 )
@@ -2841,6 +2848,10 @@ def compute_aggregate_us(aggregate_id: str, traceparent: str | None = None) -> N
                             storage_bucket,
                         )
 
+                    pe_model_version = resolve_runtime_model_version_from_db(
+                        session, simulation.tax_benefit_model_version_id
+                    )
+
                     with logfire.span("load_output"):
                         pe_output_dataset = PolicyEngineUSDataset(
                             name=output_dataset.name or "output",
@@ -2852,7 +2863,7 @@ def compute_aggregate_us(aggregate_id: str, traceparent: str | None = None) -> N
 
                         pe_sim = PESimulation(
                             dataset=pe_output_dataset,
-                            tax_benefit_model_version=us_latest,
+                            tax_benefit_model_version=pe_model_version,
                         )
                         pe_sim.output_dataset = pe_output_dataset
 

src/policyengine_api/runtime_versions.py

@@ -0,0 +1,110 @@
+"""Helpers for resolving the deployed PolicyEngine runtime bundle."""
+
+from importlib import import_module
+from uuid import UUID
+
+from sqlmodel import Session
+
+
+def _normalize_model_name(model_name: str) -> str:
+    return model_name.replace("_", "-").lower()
+
+
+def _load_runtime_model_version(model_name: str):
+    normalized_name = _normalize_model_name(model_name)
+
+    if normalized_name == "policyengine-uk":
+        return import_module("policyengine.tax_benefit_models.uk").uk_latest
+    if normalized_name == "policyengine-us":
+        return import_module("policyengine.tax_benefit_models.us").us_latest
+
+    raise ValueError(f"Unsupported tax-benefit model '{model_name}'")
+
+
+def resolve_runtime_model_version_from_db(
+    session: Session,
+    tax_benefit_model_version_id: UUID,
+):
+    """Resolve the deployed policyengine model version for a stored DB row.
+
+    The current deployment only has one runtime bundle per country. If the
+    stored DB version does not match the deployed runtime bundle, fail clearly
+    instead of silently executing against `*_latest`.
+    """
+    from policyengine_api.models import TaxBenefitModel, TaxBenefitModelVersion
+
+    db_version = session.get(TaxBenefitModelVersion, tax_benefit_model_version_id)
+    if db_version is None:
+        raise ValueError(
+            f"Tax-benefit model version {tax_benefit_model_version_id} not found"
+        )
+
+    db_model = session.get(TaxBenefitModel, db_version.model_id)
+    if db_model is None:
+        raise ValueError(f"Tax-benefit model {db_version.model_id} not found")
+
+    runtime_model_version = _load_runtime_model_version(db_model.name)
+    runtime_version = getattr(runtime_model_version, "version", None)
+
+    if runtime_version != db_version.version:
+        raise ValueError(
+            "Stored tax-benefit model version "
+            f"{db_model.name}@{db_version.version} does not match the deployed "
+            f"runtime bundle {db_model.name}@{runtime_version}. "
+            "Re-seed this environment with the deployed bundle or re-run the "
+            "analysis against the currently deployed version."
+        )
+
+    return runtime_model_version
+
+
+def _resolve_runtime_model_entry_from_db(
+    session: Session,
+    tax_benefit_model_version_id: UUID,
+):
+    from policyengine_api.models import TaxBenefitModel, TaxBenefitModelVersion
+
+    db_version = session.get(TaxBenefitModelVersion, tax_benefit_model_version_id)
+    if db_version is None:
+        raise ValueError(
+            f"Tax-benefit model version {tax_benefit_model_version_id} not found"
+        )
+
+    db_model = session.get(TaxBenefitModel, db_version.model_id)
+    if db_model is None:
+        raise ValueError(f"Tax-benefit model {db_version.model_id} not found")
+
+    runtime_model_version = resolve_runtime_model_version_from_db(
+        session, tax_benefit_model_version_id
+    )
+    return _normalize_model_name(db_model.name), runtime_model_version
+
+
+def resolve_shared_runtime_model_version_from_db(
+    session: Session,
+    *tax_benefit_model_version_ids: UUID,
+):
+    """Resolve one deployed runtime model version shared across DB version rows."""
+    if not tax_benefit_model_version_ids:
+        raise ValueError("At least one tax-benefit model version ID is required")
+
+    resolved_entries = [
+        _resolve_runtime_model_entry_from_db(session, version_id)
+        for version_id in tax_benefit_model_version_ids
+    ]
+    first_model_name, first_version = resolved_entries[0]
+
+    for model_name, runtime_model_version in resolved_entries[1:]:
+        if model_name != first_model_name:
+            raise ValueError(
+                "All simulations in a comparison must use the same tax-benefit "
+                f"model. Got {first_model_name} and {model_name}."
+            )
+        if runtime_model_version.version != first_version.version:
+            raise ValueError(
+                "All simulations in a comparison must use the same deployed "
+                f"runtime bundle. Got {first_version.version} and "
+                f"{runtime_model_version.version}."
+            )
+
+    return first_version