Harden long-run production metadata (#997)

Author: MaxGhenis

changelog.d/longrun-policyengine-build-metadata.fixed.md

@@ -0,0 +1 @@
+Stamp long-run H5 metadata with exact policyengine-us build provenance.

policyengine_us_data/datasets/cps/long_term/calibration_artifacts.py

@@ -2,12 +2,15 @@
 
 from datetime import datetime, timezone
 import hashlib
-from importlib import metadata as importlib_metadata
 import json
 from pathlib import Path
-import subprocess
 from typing import Any
 
+from policyengine_us_data.utils.policyengine import (
+    PolicyEngineUSBuildInfo,
+    get_policyengine_us_build_info,
+)
+
 try:
     from .calibration_profiles import (
         classify_calibration_quality,
@@ -27,6 +30,38 @@
 SUPPORT_AUGMENTATION_REPORT_FILENAME = "support_augmentation_report.json"
 
 
+def _coerce_policyengine_us_metadata(
+    policyengine_us: PolicyEngineUSBuildInfo | dict[str, Any] | None,
+) -> dict[str, Any]:
+    if policyengine_us is None:
+        return get_policyengine_us_build_info().to_metadata_dict()
+    if isinstance(policyengine_us, PolicyEngineUSBuildInfo):
+        return policyengine_us.to_metadata_dict()
+    coerced = json.loads(json.dumps(policyengine_us))
+    git_head = coerced.get("git_head")
+    if git_head is not None:
+        coerced.setdefault("git_commit", git_head)
+        coerced.setdefault("commit_id", git_head)
+        coerced.setdefault(
+            "direct_url",
+            {"vcs_info": {"commit_id": git_head, "vcs": "git"}},
+        )
+    return coerced
+
+
+def _metadata_policyengine_us(
+    metadata_path: str | Path,
+) -> dict[str, Any] | None:
+    path = Path(metadata_path)
+    if not path.exists():
+        return None
+    metadata = json.loads(path.read_text(encoding="utf-8"))
+    policyengine_us = metadata.get("policyengine_us")
+    if policyengine_us is None:
+        return None
+    return json.loads(json.dumps(policyengine_us))
+
+
 def metadata_path_for(h5_path: str | Path) -> Path:
     return Path(f"{Path(h5_path)}.metadata.json")
 
@@ -39,55 +74,8 @@ def _sha256_file(path: Path) -> str:
     return hashlib.sha256(path.read_bytes()).hexdigest()
 
 
-def _find_git_repo_root(path: Path) -> Path | None:
-    current = path if path.is_dir() else path.parent
-    for candidate in (current, *current.parents):
-        if (candidate / ".git").exists():
-            return candidate
-    return None
-
-
 def capture_policyengine_us_provenance() -> dict[str, Any]:
-    import policyengine_us
-
-    package_file = Path(policyengine_us.__file__).resolve()
-    version = getattr(policyengine_us, "__version__", None)
-    if version is None:
-        try:
-            version = importlib_metadata.version("policyengine-us")
-        except importlib_metadata.PackageNotFoundError:
-            version = None
-    provenance: dict[str, Any] = {
-        "package_file": str(package_file),
-        "package_file_sha256": _sha256_file(package_file),
-        "package_mtime_ns": package_file.stat().st_mtime_ns,
-        "package_size": package_file.stat().st_size,
-        "version": version,
-    }
-    repo_root = _find_git_repo_root(package_file)
-    if repo_root is None:
-        return provenance
-
-    provenance["repo_root"] = str(repo_root)
-    head = subprocess.run(
-        ["git", "rev-parse", "HEAD"],
-        cwd=repo_root,
-        check=False,
-        capture_output=True,
-        text=True,
-    )
-    if head.returncode == 0:
-        provenance["git_head"] = head.stdout.strip()
-    status = subprocess.run(
-        ["git", "status", "--porcelain=v1"],
-        cwd=repo_root,
-        check=False,
-        capture_output=True,
-        text=True,
-    )
-    if status.returncode == 0:
-        provenance["git_dirty"] = bool(status.stdout.strip())
-    return provenance
+    return get_policyengine_us_build_info().to_metadata_dict()
 
 
 def _resolve_base_dataset_path(base_dataset_path: str) -> Path | None:
@@ -237,7 +225,7 @@ def write_year_metadata(
     target_source: dict[str, Any] | None = None,
     tax_assumption: dict[str, Any] | None = None,
     support_augmentation: dict[str, Any] | None = None,
-    policyengine_us: dict[str, Any] | None = None,
+    policyengine_us: PolicyEngineUSBuildInfo | dict[str, Any] | None = None,
     base_dataset_snapshot: dict[str, Any] | None = None,
 ) -> Path:
     metadata = {
@@ -246,15 +234,14 @@ def write_year_metadata(
         "base_dataset_path": base_dataset_path,
         "profile": profile,
         "calibration_audit": calibration_audit,
+        "policyengine_us": _coerce_policyengine_us_metadata(policyengine_us),
     }
     if target_source is not None:
         metadata["target_source"] = target_source
     if tax_assumption is not None:
         metadata["tax_assumption"] = tax_assumption
     if support_augmentation is not None:
         metadata["support_augmentation"] = support_augmentation
-    if policyengine_us is not None:
-        metadata["policyengine_us"] = policyengine_us
     if base_dataset_snapshot is not None:
         metadata["base_dataset_snapshot"] = base_dataset_snapshot
     metadata = normalize_metadata(metadata)
@@ -294,7 +281,7 @@ def update_dataset_manifest(
     target_source: dict[str, Any] | None = None,
     tax_assumption: dict[str, Any] | None = None,
     support_augmentation: dict[str, Any] | None = None,
-    policyengine_us: dict[str, Any] | None = None,
+    policyengine_us: PolicyEngineUSBuildInfo | dict[str, Any] | None = None,
     base_dataset_snapshot: dict[str, Any] | None = None,
 ) -> Path:
     output_dir = Path(output_dir)
@@ -303,7 +290,13 @@ def update_dataset_manifest(
     target_source = _json_clone(target_source)
     tax_assumption = _json_clone(tax_assumption)
     support_augmentation = _json_clone(support_augmentation)
-    policyengine_us = _json_clone(policyengine_us)
+    policyengine_us = (
+        _coerce_policyengine_us_metadata(policyengine_us)
+        if policyengine_us is not None
+        else _metadata_policyengine_us(metadata_path)
+    )
+    if policyengine_us is None:
+        policyengine_us = get_policyengine_us_build_info().to_metadata_dict()
     base_dataset_snapshot = _json_clone(base_dataset_snapshot)
 
     if manifest_path.exists():
@@ -402,6 +395,7 @@ def update_dataset_manifest(
         "negative_weight_household_pct": calibration_audit.get(
             "negative_weight_household_pct"
         ),
+        "policyengine_us_version": policyengine_us.get("version"),
         "validation_passed": calibration_audit.get("validation_passed"),
         "validation_issue_count": len(calibration_audit.get("validation_issues", [])),
     }

policyengine_us_data/datasets/cps/long_term/prototype_synthetic_2100_support.py

@@ -110,6 +110,45 @@
 )
 PAYROLL_UPRATING_FACTOR_COLUMN = "__pe_payroll_uprating_factor"
 SS_UPRATING_FACTOR_COLUMN = "__pe_ss_uprating_factor"
+NON_TARGET_CLONE_INCOME_COMPONENTS = (
+    "alimony_income",
+    "disability_benefits",
+    "estate_income",
+    "farm_income",
+    "farm_operations_income",
+    "farm_rent_income",
+    "keogh_distributions",
+    "long_term_capital_gains_before_response",
+    "long_term_capital_gains_on_collectibles",
+    "miscellaneous_income",
+    "non_qualified_dividend_income",
+    "non_sch_d_capital_gains",
+    "partnership_s_corp_income",
+    "partnership_se_income",
+    "qualified_bdc_income",
+    "qualified_dividend_income",
+    "qualified_reit_and_ptp_income",
+    "rental_income",
+    "salt_refund_income",
+    "short_term_capital_gains",
+    "strike_benefits",
+    "tax_exempt_401k_distributions",
+    "tax_exempt_403b_distributions",
+    "tax_exempt_interest_income",
+    "tax_exempt_ira_distributions",
+    "tax_exempt_private_pension_income",
+    "tax_exempt_sep_distributions",
+    "taxable_401k_distributions",
+    "taxable_403b_distributions",
+    "taxable_interest_income",
+    "taxable_ira_distributions",
+    "taxable_private_pension_income",
+    "taxable_sep_distributions",
+    "unemployment_compensation",
+    "unrecaptured_section_1250_gain",
+    "veterans_benefits",
+    "workers_compensation",
+)
 
 
 @dataclass(frozen=True)
@@ -685,14 +724,24 @@ def _scale_levels(levels: list[float], scale: float) -> list[float]:
 
 
 @lru_cache(maxsize=None)
-def load_policyengine_social_security_cap(year: int) -> float:
+def _load_policyengine_social_security_cap_without_reform(year: int) -> float:
     sim = Microsimulation(dataset=DEFAULT_DATASET)
     return validate_projected_social_security_cap(
         sim.tax_benefit_system.parameters,
         year,
     )
 
 
+def load_policyengine_social_security_cap(year: int, *, reform=None) -> float:
+    if reform is None:
+        return _load_policyengine_social_security_cap_without_reform(year)
+    sim = Microsimulation(dataset=DEFAULT_DATASET, reform=reform)
+    return validate_projected_social_security_cap(
+        sim.tax_benefit_system.parameters,
+        year,
+    )
+
+
 def allocate_taxable_payroll_wages(
     total_taxable_payroll: float,
     payroll_split: tuple[float, float],
@@ -1868,6 +1917,21 @@ def _zero_period_columns(
     return available_columns
 
 
+def _zero_clone_non_target_income(
+    cloned: pd.DataFrame,
+    *,
+    base_year: int,
+) -> pd.DataFrame:
+    columns = [
+        _period_column(component, base_year)
+        for component in NON_TARGET_CLONE_INCOME_COMPONENTS
+        if _period_column(component, base_year) in cloned.columns
+    ]
+    if columns:
+        cloned.loc[:, columns] = 0.0
+    return cloned
+
+
 def _target_base_total_for_row(
     row: pd.Series,
     *,
@@ -1964,6 +2028,7 @@ def _clone_tax_unit_rows_to_target(
         _period_column(component, base_year) for component in SS_COMPONENTS
     )
     qbi_col = _period_column("w2_wages_from_qualified_business", base_year)
+    cloned = _zero_clone_non_target_income(cloned, base_year=base_year)
 
     target_head_payroll = _target_base_total_for_row(
         cloned.loc[head_idx],
@@ -2351,7 +2416,10 @@ def build_donor_backed_augmented_input_dataframe(
         ss_scale=ss_scale,
         earnings_scale=earnings_scale,
     )
-    payroll_cap = load_policyengine_social_security_cap(target_year)
+    payroll_cap = load_policyengine_social_security_cap(
+        target_year,
+        reform=reform,
+    )
     candidates = generate_synthetic_candidates(
         pools,
         payroll_cap=payroll_cap,
@@ -2499,7 +2567,10 @@ def build_role_composite_augmented_input_dataframe(
         ss_scale=ss_scale,
         earnings_scale=earnings_scale,
     )
-    payroll_cap = load_policyengine_social_security_cap(target_year)
+    payroll_cap = load_policyengine_social_security_cap(
+        target_year,
+        reform=reform,
+    )
     candidates = generate_synthetic_candidates(
         pools,
         payroll_cap=payroll_cap,

policyengine_us_data/datasets/cps/long_term/run_long_term_production.py

@@ -15,6 +15,10 @@
     parse_years,
 )
 from policyengine_us_data.utils.data_upload import upload_to_staging_hf
+from policyengine_us_data.utils.policyengine import (
+    PolicyEngineUSBuildInfo,
+    assert_locked_policyengine_us_version,
+)
 from policyengine_us_data.utils.run_context import resolve_run_id, staging_prefix
 
 
@@ -176,6 +180,7 @@ def write_manifest(
     years: list[int],
     run_id: str,
     source_sha: str,
+    policyengine_us_build: PolicyEngineUSBuildInfo,
     artifacts: list[Path],
 ) -> Path:
     manifest_path = output_dir / "long_run_production_manifest.json"
@@ -197,6 +202,7 @@ def write_manifest(
             "policyengine-us": _package_version("policyengine-us"),
             "policyengine-core": _package_version("policyengine-core"),
         },
+        "policyengine_us": policyengine_us_build.to_metadata_dict(),
         "projection": {
             "years_spec": args.years,
             "years": years,
@@ -329,6 +335,7 @@ def main() -> int:
     source_sha = args.source_sha or os.environ.get("GITHUB_SHA", "") or _git_sha()
 
     command = build_projection_command(args, output_dir)
+    policyengine_us_build = assert_locked_policyengine_us_version()
     print("Running long-run projection command:")
     print(" ".join(command))
     subprocess.run(command, check=True)
@@ -341,6 +348,7 @@ def main() -> int:
         years=years,
         run_id=run_id,
         source_sha=source_sha,
+        policyengine_us_build=policyengine_us_build,
         artifacts=artifacts,
     )
     artifacts = collect_artifacts(output_dir, args.artifact_prefix)