diff --git a/changelog.d/versioned-trace-export.added.md b/changelog.d/versioned-trace-export.added.md new file mode 100644 index 00000000..77643ee9 --- /dev/null +++ b/changelog.d/versioned-trace-export.added.md @@ -0,0 +1 @@ +Added `FullTracer.to_trace()` and `Simulation.to_trace()` for versioned computation trace export (`policyengine.trace.v1`). diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py index b7f52591..04d554ca 100644 --- a/policyengine_core/simulations/simulation.py +++ b/policyengine_core/simulations/simulation.py @@ -528,6 +528,52 @@ def trace(self, trace: SimpleTracer) -> None: else: self.tracer = SimpleTracer() + def to_trace( + self, + format: str = "policyengine.trace.v1", + *, + model: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Export a versioned computation trace for this simulation. + + Requires ``trace=True`` (or assigning ``simulation.trace = True``) + so a :class:`~policyengine_core.tracers.FullTracer` is installed. + + :param format: Trace document format identifier. + :param model: Optional country-package metadata override. When omitted, + a best-effort package name/version is taken from the tax-benefit + system module when available. + :returns: JSON-compatible trace document (see + :meth:`~policyengine_core.tracers.FullTracer.to_trace`). + """ + if not isinstance(self.tracer, FullTracer): + raise ValueError( + "Computation traces require FullTracer. " + "Create the simulation with trace=True, or set simulation.trace = True " + "before calculating, then call to_trace()." + ) + if model is None: + model = self._model_metadata_for_trace() + return self.tracer.to_trace(format=format, model=model) + + def _model_metadata_for_trace(self) -> Optional[Dict[str, Any]]: + tax_benefit_system = getattr(self, "tax_benefit_system", None) + if tax_benefit_system is None: + return None + module = type(tax_benefit_system).__module__ + package_name = module.split(".", 1)[0] if module else None + if not package_name or package_name == "policyengine_core": + return None + metadata: Dict[str, Any] = {"package": package_name.replace("_", "-")} + try: + import importlib.metadata + + distribution_name = package_name.replace("_", "-") + metadata["version"] = importlib.metadata.version(distribution_name) + except Exception: + pass + return metadata + def link_to_entities_instances(self) -> None: for _key, entity_instance in self.populations.items(): entity_instance.simulation = self diff --git a/policyengine_core/tracers/full_tracer.py b/policyengine_core/tracers/full_tracer.py index 3cab802e..4afb10cb 100644 --- a/policyengine_core/tracers/full_tracer.py +++ b/policyengine_core/tracers/full_tracer.py @@ -1,8 +1,9 @@ from __future__ import annotations +import re import time import typing -from typing import Dict, Iterator, List, Optional, Union +from typing import Any, Dict, Iterator, List, Optional, Union from .. import tracers @@ -14,6 +15,25 @@ Stack = List[Dict[str, Union[str, Period]]] +TRACE_FORMAT_V1 = "policyengine.trace.v1" +# Matches FlatTrace.key(): ``name`` +_TRACE_KEY_RE = re.compile( + r"^(?P.+)<(?P[^<>]+), \((?P[^()]*)\)>$" +) + + +def parse_trace_key(key: str) -> Dict[str, str]: + """Parse a serialized flat-trace node key into name/period/branch.""" + match = _TRACE_KEY_RE.match(key) + if not match: + raise ValueError(f"Invalid PolicyEngine trace key: {key!r}") + return { + "name": match.group("name"), + "period": match.group("period").strip(), + "branch": match.group("branch").strip(), + } + + class FullTracer: _simple_tracer: tracers.SimpleTracer _trees: list @@ -172,3 +192,130 @@ def _browse_node(node): for node in self._trees: yield from _browse_node(node) + + def to_trace( + self, + format: str = TRACE_FORMAT_V1, + *, + model: Optional[Dict[str, Any]] = None, + ) -> Dict[str, Any]: + """Export a versioned, JSON-compatible computation trace. + + Node identity fields (``variable`` / ``period`` / ``branch``) are taken + from structured :class:`~policyengine_core.tracers.TraceNode` attributes + via :meth:`browse_trace`, not by regex-parsing serialized flat-trace + keys. :func:`parse_trace_key` remains available for callers that only + have a flat key string (for example parameter maps without a node). + + **v1 intentionally omits** entity/count on ``calculation`` roots and + source/version on ``parameters``. Those were listed as optional + "minimum useful fields" in the provenance issue and can land in a later + format version without breaking ``policyengine.trace.v1`` consumers. + + This document is *per-computation* provenance. It is complementary to + PolicyEngine's TRACE Transparent Research Object (TRO) surface in + ``policyengine`` (JSON-LD under ``https://policyengine.org/trace/0.1#``), + which pins release/composition artifacts. A future TRO may embed or + reference a ``policyengine.trace.v1`` payload as its runtime trace; the + format id and TRO namespace are intentionally distinct. + + Values are full serialized vectors (household-audit friendly). Optional + summarization for population-scale traces is deferred. + + :param format: Trace document format identifier. Currently only + ``policyengine.trace.v1`` is supported. + :param model: Optional country-package / model metadata + (for example ``{"package": "policyengine-us", "version": "…"}``). + :returns: A plain ``dict`` ready for ``json.dumps``. + :raises ValueError: If ``format`` is not a supported version string. + """ + if format != TRACE_FORMAT_V1: + raise ValueError( + f"Unsupported trace format {format!r}. " + f"Supported formats: {TRACE_FORMAT_V1!r}" + ) + + engine = self._engine_metadata() + flat_trace = self.flat_trace + nodes: List[Dict[str, Any]] = [] + parameters: Dict[str, Any] = {} + seen_node_ids: set[str] = set() + + for tree_node in self.browse_trace(): + node_id = flat_trace.key(tree_node) + # Cache reads can revisit the same calculation; keep the first + # structured record (same non-overwriting rule as FlatTrace). + if node_id in seen_node_ids: + continue + seen_node_ids.add(node_id) + + parameter_map = { + flat_trace.key(parameter): flat_trace.serialize(parameter.value) + for parameter in tree_node.parameters + } + node: Dict[str, Any] = { + "id": node_id, + "variable": tree_node.name, + "period": str(tree_node.period), + "branch": tree_node.branch_name, + "dependencies": [ + flat_trace.key(child) for child in tree_node.children + ], + "parameters": parameter_map, + "value": flat_trace.serialize(tree_node.value), + "calculation_time": tree_node.calculation_time(), + "formula_time": tree_node.formula_time(), + } + nodes.append(node) + + for parameter in tree_node.parameters: + parameter_id = flat_trace.key(parameter) + if parameter_id in parameters: + continue + parameters[parameter_id] = { + "id": parameter_id, + "name": parameter.name, + "instant": str(parameter.period), + "branch": parameter.branch_name, + "value": flat_trace.serialize(parameter.value), + } + + document: Dict[str, Any] = { + "format": format, + "engine": engine, + "calculation": { + "roots": [ + { + "id": flat_trace.key(tree), + "variable": tree.name, + "period": str(tree.period), + "branch": tree.branch_name, + } + for tree in self._trees + ], + }, + "nodes": nodes, + "parameters": parameters, + } + if model is not None: + document["model"] = model + return document + + @staticmethod + def _engine_metadata() -> Dict[str, Any]: + try: + from policyengine_core.build_metadata import get_runtime_metadata + + metadata = get_runtime_metadata() + engine: Dict[str, Any] = { + "package": metadata.get("name", "policyengine-core"), + "version": metadata.get("version", "unknown"), + } + if metadata.get("git_sha"): + engine["git_sha"] = metadata["git_sha"] + return engine + except Exception: + return { + "package": "policyengine-core", + "version": "unknown", + } diff --git a/tests/core/test_tracers.py b/tests/core/test_tracers.py index e653f39b..67430cfd 100644 --- a/tests/core/test_tracers.py +++ b/tests/core/test_tracers.py @@ -17,6 +17,7 @@ TraceNode, TracingParameterNodeAtInstant, ) +from policyengine_core.tracers.full_tracer import parse_trace_key from .parameters_fancy_indexing.test_fancy_indexing import parameters @@ -529,3 +530,150 @@ def test_browse_trace(): browsed_nodes = [node.name for node in tracer.browse_trace()] assert browsed_nodes == ["B", "C", "D", "E", "F"] + + +def test_to_trace_versioned_export(): + tracer = FullTracer() + tracer.record_calculation_start("income_tax", 2017) + tracer.record_calculation_start("salary", 2017) + tracer.record_calculation_result(np.asarray([2000])) + tracer.record_calculation_end() + tracer.record_parameter_access("taxes.rate", "2017-01-01", "default", 0.15) + tracer.record_calculation_result(np.asarray([300])) + tracer.record_calculation_end() + + trace = tracer.to_trace() + + assert trace["format"] == "policyengine.trace.v1" + assert trace["engine"]["package"] == "policyengine-core" + assert "version" in trace["engine"] + assert trace["calculation"]["roots"] == [ + { + "id": "income_tax<2017, (default)>", + "variable": "income_tax", + "period": "2017", + "branch": "default", + } + ] + + nodes_by_id = {node["id"]: node for node in trace["nodes"]} + root = nodes_by_id["income_tax<2017, (default)>"] + assert root["variable"] == "income_tax" + assert root["period"] == "2017" + assert root["branch"] == "default" + assert root["dependencies"] == ["salary<2017, (default)>"] + assert root["value"] == [300] + assert "taxes.rate<2017-01-01, (default)>" in root["parameters"] + assert root["parameters"]["taxes.rate<2017-01-01, (default)>"] == approx(0.15) + assert "calculation_time" in root + assert "formula_time" in root + + assert nodes_by_id["salary<2017, (default)>"]["value"] == [2000] + assert "taxes.rate<2017-01-01, (default)>" in trace["parameters"] + assert ( + trace["parameters"]["taxes.rate<2017-01-01, (default)>"]["name"] == "taxes.rate" + ) + + +def test_to_trace_rejects_unknown_format(): + tracer = FullTracer() + tracer.record_calculation_start("salary", 2017) + tracer.record_calculation_end() + + with raises(ValueError, match="Unsupported trace format"): + tracer.to_trace(format="policyengine.trace.v0") + + +def test_to_trace_includes_optional_model_metadata(): + tracer = FullTracer() + tracer.record_calculation_start("salary", 2017) + tracer.record_calculation_end() + + trace = tracer.to_trace(model={"package": "policyengine-us", "version": "0.0.0"}) + assert trace["model"] == {"package": "policyengine-us", "version": "0.0.0"} + + +def test_simulation_to_trace_requires_full_tracer(): + simulation = StubSimulation() + simulation.tracer = SimpleTracer() + + with raises(ValueError, match="FullTracer"): + simulation.to_trace() + + +def test_simulation_to_trace_exports_full_tracer_document(): + simulation = StubSimulation() + simulation.tracer = FullTracer() + simulation.tax_benefit_system = None + simulation.tracer.record_calculation_start("salary", 2017) + simulation.tracer.record_calculation_result(np.asarray([1000])) + simulation.tracer.record_calculation_end() + + trace = simulation.to_trace(model={"package": "policyengine-us", "version": "1.0.0"}) + + assert trace["format"] == "policyengine.trace.v1" + assert trace["model"] == {"package": "policyengine-us", "version": "1.0.0"} + assert trace["calculation"]["roots"][0]["variable"] == "salary" + assert trace["nodes"][0]["value"] == [1000] + + +def test_simulation_to_trace_infers_model_metadata_from_tax_benefit_system(): + class FakeCountrySystem: + pass + + FakeCountrySystem.__module__ = "policyengine_us.system" + + simulation = StubSimulation() + simulation.tracer = FullTracer() + simulation.tax_benefit_system = FakeCountrySystem() + simulation.tracer.record_calculation_start("salary", 2017) + simulation.tracer.record_calculation_end() + + trace = simulation.to_trace() + + assert trace["model"]["package"] == "policyengine-us" + + +def test_simulation_to_trace_skips_core_package_as_model_metadata(): + class CoreLikeSystem: + pass + + CoreLikeSystem.__module__ = "policyengine_core.country_template" + + simulation = StubSimulation() + simulation.tracer = FullTracer() + simulation.tax_benefit_system = CoreLikeSystem() + simulation.tracer.record_calculation_start("salary", 2017) + simulation.tracer.record_calculation_end() + + trace = simulation.to_trace() + assert "model" not in trace + + +def test_to_trace_reads_structured_fields_from_trace_nodes(): + """Nodes/parameters come from TraceNode attributes, not key regex parsing.""" + tracer = FullTracer() + tracer.record_calculation_start("income_tax", 2017, branch_name="reform") + tracer.record_parameter_access("taxes.rate", "2017-01-01", "reform", 0.2) + tracer.record_calculation_result(np.asarray([100])) + tracer.record_calculation_end() + + trace = tracer.to_trace() + root = trace["nodes"][0] + assert root["variable"] == "income_tax" + assert root["period"] == "2017" + assert root["branch"] == "reform" + + parameter = trace["parameters"]["taxes.rate<2017-01-01, (reform)>"] + assert parameter["name"] == "taxes.rate" + assert parameter["instant"] == "2017-01-01" + assert parameter["branch"] == "reform" + assert parameter["value"] == approx(0.2) + + +def test_parse_trace_key_still_available_for_flat_keys(): + parsed = parse_trace_key("salary<2017, (default)>") + assert parsed == {"name": "salary", "period": "2017", "branch": "default"} + + with raises(ValueError, match="Invalid PolicyEngine trace key"): + parse_trace_key("not-a-valid-trace-key")