initial code
This commit is contained in:
@@ -0,0 +1,727 @@
|
||||
"""Build a private, point-in-time inspection feature mart with DuckDB.
|
||||
|
||||
Inputs are either a contiguous set of bounded monthly CSV.gz exports or the
|
||||
single complete-history development sample. Every input must have the manifest
|
||||
written by the corresponding secure exporter.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import gzip
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple
|
||||
|
||||
import duckdb
|
||||
|
||||
|
||||
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
||||
PRIVATE_DATA_ROOT = (PROJECT_ROOT / "data/private").resolve()
|
||||
DEFAULT_INPUT = PRIVATE_DATA_ROOT / "inspection_batches"
|
||||
DEFAULT_DATABASE = PRIVATE_DATA_ROOT / "warehouse/vehicle_health.duckdb"
|
||||
DEFAULT_OUTPUT = PRIVATE_DATA_ROOT / "marts/inspection_feature_mart.parquet"
|
||||
SQL_ROOT = PROJECT_ROOT / "sql/local"
|
||||
|
||||
EXPECTED_COLUMNS = (
|
||||
"internal_event_id",
|
||||
"vehicle_token",
|
||||
"vehicle_bucket",
|
||||
"event_ts",
|
||||
"source_era",
|
||||
"public_county",
|
||||
"canonical_outcome",
|
||||
"outcome_label_source",
|
||||
"program_type",
|
||||
"test_type",
|
||||
"observed_make",
|
||||
"observed_model",
|
||||
"observed_model_year",
|
||||
)
|
||||
BOUNDED_KIND = "bounded_batch"
|
||||
DEVELOPMENT_KIND = "vehicle_history_development_sample"
|
||||
HEX_64 = re.compile(r"^[0-9a-f]{64}$")
|
||||
FINGERPRINT = re.compile(r"^[0-9a-f]{16,64}$")
|
||||
MEMORY_LIMIT = re.compile(r"^[1-9][0-9]*(?:MB|GB|TB)$", re.IGNORECASE)
|
||||
|
||||
SQL_FILES = (
|
||||
"20_stage_events.sql",
|
||||
"21_build_episodes.sql",
|
||||
"22_build_features.sql",
|
||||
"23_validate_mart.sql",
|
||||
)
|
||||
|
||||
CREATE_STAGING_SQL = """
|
||||
CREATE TABLE stg_events (
|
||||
batch_file VARCHAR NOT NULL,
|
||||
batch_kind VARCHAR NOT NULL,
|
||||
batch_start TIMESTAMP NOT NULL,
|
||||
batch_end TIMESTAMP NOT NULL,
|
||||
internal_event_id VARCHAR,
|
||||
vehicle_token VARCHAR,
|
||||
vehicle_bucket VARCHAR,
|
||||
event_ts VARCHAR,
|
||||
source_era VARCHAR,
|
||||
public_county VARCHAR,
|
||||
canonical_outcome VARCHAR,
|
||||
outcome_label_source VARCHAR,
|
||||
program_type VARCHAR,
|
||||
test_type VARCHAR,
|
||||
observed_make VARCHAR,
|
||||
observed_model VARCHAR,
|
||||
observed_model_year VARCHAR
|
||||
)
|
||||
"""
|
||||
|
||||
# The dialect and every input column are explicit. In particular, do not let a
|
||||
# sample without quoted values convince DuckDB that quote handling is disabled.
|
||||
INGEST_SQL = """
|
||||
INSERT INTO stg_events
|
||||
SELECT
|
||||
? AS batch_file,
|
||||
? AS batch_kind,
|
||||
?::TIMESTAMP AS batch_start,
|
||||
?::TIMESTAMP AS batch_end,
|
||||
source.internal_event_id,
|
||||
source.vehicle_token,
|
||||
source.vehicle_bucket,
|
||||
source.event_ts,
|
||||
source.source_era,
|
||||
source.public_county,
|
||||
source.canonical_outcome,
|
||||
source.outcome_label_source,
|
||||
source.program_type,
|
||||
source.test_type,
|
||||
source.observed_make,
|
||||
source.observed_model,
|
||||
source.observed_model_year
|
||||
FROM read_csv(
|
||||
?,
|
||||
header = true,
|
||||
auto_detect = false,
|
||||
delim = ',',
|
||||
quote = '"',
|
||||
escape = '"',
|
||||
nullstr = '',
|
||||
strict_mode = true,
|
||||
columns = {
|
||||
'internal_event_id': 'VARCHAR',
|
||||
'vehicle_token': 'VARCHAR',
|
||||
'vehicle_bucket': 'VARCHAR',
|
||||
'event_ts': 'VARCHAR',
|
||||
'source_era': 'VARCHAR',
|
||||
'public_county': 'VARCHAR',
|
||||
'canonical_outcome': 'VARCHAR',
|
||||
'outcome_label_source': 'VARCHAR',
|
||||
'program_type': 'VARCHAR',
|
||||
'test_type': 'VARCHAR',
|
||||
'observed_make': 'VARCHAR',
|
||||
'observed_model': 'VARCHAR',
|
||||
'observed_model_year': 'VARCHAR'
|
||||
}
|
||||
) AS source
|
||||
"""
|
||||
|
||||
|
||||
class MartBuildError(RuntimeError):
|
||||
"""Raised when private input or a mart invariant fails closed."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class InputBatch:
|
||||
path: Path
|
||||
manifest_path: Path
|
||||
export_kind: str
|
||||
start: datetime
|
||||
end: datetime
|
||||
rows: int
|
||||
query_version: str
|
||||
query_sha256: str
|
||||
key_version: str
|
||||
key_fingerprint: str
|
||||
compressed_sha256: str
|
||||
compressed_bytes: int
|
||||
manifest: Mapping[str, Any]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BuildConfig:
|
||||
input_path: Path = DEFAULT_INPUT
|
||||
database_path: Path = DEFAULT_DATABASE
|
||||
output_path: Path = DEFAULT_OUTPUT
|
||||
episode_gap_days: int = 30
|
||||
memory_limit: str = "4GB"
|
||||
threads: int = max(1, min(4, os.cpu_count() or 1))
|
||||
allow_gaps: bool = False
|
||||
allow_external_output: bool = False
|
||||
overwrite: bool = False
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BuildSummary:
|
||||
database_path: Path
|
||||
output_path: Path
|
||||
manifest_path: Path
|
||||
input_rows: int
|
||||
clean_events: int
|
||||
episodes: int
|
||||
mart_rows: int
|
||||
eligible_rows: int
|
||||
source_data_kind: str
|
||||
population_estimate_allowed: bool
|
||||
|
||||
|
||||
def file_sha256(path: Path) -> str:
|
||||
digest = hashlib.sha256()
|
||||
with path.open("rb") as handle:
|
||||
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
|
||||
digest.update(chunk)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def parse_manifest_timestamp(value: object, field: str) -> datetime:
|
||||
if not isinstance(value, str):
|
||||
raise MartBuildError(f"Manifest field {field!r} must be an ISO timestamp")
|
||||
encoded = value[:-1] + "+00:00" if value.endswith("Z") else value
|
||||
try:
|
||||
parsed = datetime.fromisoformat(encoded)
|
||||
except ValueError as exc:
|
||||
raise MartBuildError(f"Manifest field {field!r} is not ISO-8601") from exc
|
||||
if parsed.tzinfo is not None:
|
||||
parsed = parsed.astimezone(timezone.utc).replace(tzinfo=None)
|
||||
return parsed
|
||||
|
||||
|
||||
def _required_text(manifest: Mapping[str, Any], field: str) -> str:
|
||||
value = manifest.get(field)
|
||||
if not isinstance(value, str) or not value.strip():
|
||||
raise MartBuildError(f"Manifest field {field!r} must be non-empty text")
|
||||
return value.strip()
|
||||
|
||||
|
||||
def _required_nonnegative_int(manifest: Mapping[str, Any], field: str) -> int:
|
||||
value = manifest.get(field)
|
||||
if isinstance(value, bool) or not isinstance(value, int) or value < 0:
|
||||
raise MartBuildError(f"Manifest field {field!r} must be a nonnegative integer")
|
||||
return value
|
||||
|
||||
|
||||
def _read_header(path: Path) -> Tuple[str, ...]:
|
||||
try:
|
||||
with gzip.open(path, "rt", encoding="utf-8", newline="") as handle:
|
||||
row = next(csv.reader(handle), None)
|
||||
except (OSError, UnicodeError, csv.Error) as exc:
|
||||
raise MartBuildError(f"Cannot read gzip CSV header: {path}") from exc
|
||||
if row is None:
|
||||
raise MartBuildError(f"Input gzip CSV is empty: {path}")
|
||||
return tuple(row)
|
||||
|
||||
|
||||
def _load_batch(path: Path) -> InputBatch:
|
||||
manifest_path = path.with_name(path.name + ".manifest.json")
|
||||
if not manifest_path.is_file():
|
||||
raise MartBuildError(f"Missing paired manifest: {manifest_path}")
|
||||
try:
|
||||
manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
|
||||
except (OSError, UnicodeError, json.JSONDecodeError) as exc:
|
||||
raise MartBuildError(f"Cannot parse manifest: {manifest_path}") from exc
|
||||
if not isinstance(manifest, dict):
|
||||
raise MartBuildError(f"Manifest must contain a JSON object: {manifest_path}")
|
||||
|
||||
columns = manifest.get("columns")
|
||||
if not isinstance(columns, (list, tuple)) or tuple(columns) != EXPECTED_COLUMNS:
|
||||
raise MartBuildError(f"Unexpected manifest columns: {manifest_path}")
|
||||
if _read_header(path) != EXPECTED_COLUMNS:
|
||||
raise MartBuildError(f"CSV header does not match the private schema: {path}")
|
||||
|
||||
export_kind = _required_text(manifest, "export_kind")
|
||||
if export_kind not in (BOUNDED_KIND, DEVELOPMENT_KIND):
|
||||
raise MartBuildError(f"Unsupported export_kind {export_kind!r}: {manifest_path}")
|
||||
start = parse_manifest_timestamp(
|
||||
manifest.get("source_start_inclusive"), "source_start_inclusive"
|
||||
)
|
||||
end = parse_manifest_timestamp(
|
||||
manifest.get("source_end_exclusive"), "source_end_exclusive"
|
||||
)
|
||||
if end <= start:
|
||||
raise MartBuildError(f"Manifest source interval is empty or reversed: {path}")
|
||||
|
||||
rows = _required_nonnegative_int(manifest, "rows")
|
||||
compressed_bytes = _required_nonnegative_int(manifest, "compressed_file_bytes")
|
||||
actual_bytes = path.stat().st_size
|
||||
if compressed_bytes != actual_bytes:
|
||||
raise MartBuildError(
|
||||
f"Compressed byte count mismatch for {path}: "
|
||||
f"manifest={compressed_bytes}, actual={actual_bytes}"
|
||||
)
|
||||
compressed_sha256 = _required_text(manifest, "compressed_file_sha256").lower()
|
||||
if not HEX_64.fullmatch(compressed_sha256):
|
||||
raise MartBuildError(f"Invalid compressed SHA-256 in {manifest_path}")
|
||||
actual_sha256 = file_sha256(path)
|
||||
if compressed_sha256 != actual_sha256:
|
||||
raise MartBuildError(f"Compressed SHA-256 mismatch for {path}")
|
||||
|
||||
query_sha256 = _required_text(manifest, "query_sha256").lower()
|
||||
if not HEX_64.fullmatch(query_sha256):
|
||||
raise MartBuildError(f"Invalid query SHA-256 in {manifest_path}")
|
||||
key_fingerprint = _required_text(
|
||||
manifest, "vehicle_token_key_fingerprint"
|
||||
).lower()
|
||||
if not FINGERPRINT.fullmatch(key_fingerprint):
|
||||
raise MartBuildError(f"Invalid key fingerprint in {manifest_path}")
|
||||
|
||||
if export_kind == BOUNDED_KIND:
|
||||
if manifest.get("classification") != (
|
||||
"private_pseudonymized_analytical_staging"
|
||||
):
|
||||
raise MartBuildError(f"Unexpected bounded-batch classification: {path}")
|
||||
else:
|
||||
if manifest.get("classification") != (
|
||||
"private_pseudonymized_development_only"
|
||||
):
|
||||
raise MartBuildError(f"Unexpected development classification: {path}")
|
||||
if manifest.get("population_estimate_allowed") is not False:
|
||||
raise MartBuildError(
|
||||
"Development history manifest must prohibit population estimates"
|
||||
)
|
||||
|
||||
return InputBatch(
|
||||
path=path,
|
||||
manifest_path=manifest_path,
|
||||
export_kind=export_kind,
|
||||
start=start,
|
||||
end=end,
|
||||
rows=rows,
|
||||
query_version=_required_text(manifest, "query_version"),
|
||||
query_sha256=query_sha256,
|
||||
key_version=_required_text(manifest, "vehicle_token_key_version"),
|
||||
key_fingerprint=key_fingerprint,
|
||||
compressed_sha256=compressed_sha256,
|
||||
compressed_bytes=compressed_bytes,
|
||||
manifest=manifest,
|
||||
)
|
||||
|
||||
|
||||
def _discover_files(input_path: Path) -> List[Path]:
|
||||
if input_path.is_file():
|
||||
if not input_path.name.endswith(".csv.gz"):
|
||||
raise MartBuildError("A feature-mart input file must end in .csv.gz")
|
||||
return [input_path]
|
||||
if not input_path.is_dir():
|
||||
raise MartBuildError(f"Input path does not exist: {input_path}")
|
||||
paths = sorted(path for path in input_path.glob("*.csv.gz") if path.is_file())
|
||||
if not paths:
|
||||
raise MartBuildError(f"No .csv.gz inputs found in {input_path}")
|
||||
return paths
|
||||
|
||||
|
||||
def discover_and_validate_batches(config: BuildConfig) -> List[InputBatch]:
|
||||
input_path = config.input_path.resolve()
|
||||
if not config.allow_external_output and not input_path.is_relative_to(
|
||||
PRIVATE_DATA_ROOT
|
||||
):
|
||||
raise MartBuildError(f"Private inputs must stay under {PRIVATE_DATA_ROOT}")
|
||||
|
||||
batches = [_load_batch(path.resolve()) for path in _discover_files(input_path)]
|
||||
kinds = {batch.export_kind for batch in batches}
|
||||
if len(kinds) != 1:
|
||||
raise MartBuildError("Do not mix bounded batches and development samples")
|
||||
kind = next(iter(kinds))
|
||||
|
||||
key_versions = {batch.key_version for batch in batches}
|
||||
fingerprints = {batch.key_fingerprint for batch in batches}
|
||||
if len(key_versions) != 1 or len(fingerprints) != 1:
|
||||
raise MartBuildError("All inputs must use one VIN-token key and version")
|
||||
|
||||
if kind == DEVELOPMENT_KIND:
|
||||
if len(batches) != 1:
|
||||
raise MartBuildError("A development mart accepts exactly one history sample")
|
||||
return batches
|
||||
|
||||
query_versions = {batch.query_version for batch in batches}
|
||||
query_hashes = {batch.query_sha256 for batch in batches}
|
||||
if len(query_versions) != 1 or len(query_hashes) != 1:
|
||||
raise MartBuildError("All bounded batches must use one extraction query version")
|
||||
|
||||
batches.sort(key=lambda batch: (batch.start, batch.end, str(batch.path)))
|
||||
for previous, current in zip(batches, batches[1:]):
|
||||
if current.start < previous.end:
|
||||
raise MartBuildError(
|
||||
f"Source intervals overlap: {previous.path.name} and "
|
||||
f"{current.path.name}"
|
||||
)
|
||||
if current.start > previous.end and not config.allow_gaps:
|
||||
raise MartBuildError(
|
||||
f"Source interval gap between {previous.end.isoformat()} and "
|
||||
f"{current.start.isoformat()}; use --allow-gaps only for an "
|
||||
"explicitly incomplete analytical build"
|
||||
)
|
||||
return batches
|
||||
|
||||
|
||||
def _validate_config(config: BuildConfig) -> BuildConfig:
|
||||
if not 1 <= config.episode_gap_days <= 365:
|
||||
raise MartBuildError("episode_gap_days must be between 1 and 365")
|
||||
if not 1 <= config.threads <= 64:
|
||||
raise MartBuildError("threads must be between 1 and 64")
|
||||
if not MEMORY_LIMIT.fullmatch(config.memory_limit):
|
||||
raise MartBuildError("memory_limit must look like 4096MB or 4GB")
|
||||
|
||||
database_path = config.database_path.resolve()
|
||||
output_path = config.output_path.resolve()
|
||||
if output_path.suffix.lower() != ".parquet":
|
||||
raise MartBuildError("Feature-mart output must use a .parquet suffix")
|
||||
if not config.allow_external_output:
|
||||
for path in (database_path, output_path):
|
||||
if not path.is_relative_to(PRIVATE_DATA_ROOT):
|
||||
raise MartBuildError(f"Private outputs must stay under {PRIVATE_DATA_ROOT}")
|
||||
if database_path == output_path:
|
||||
raise MartBuildError("Database and Parquet output paths must differ")
|
||||
|
||||
return BuildConfig(
|
||||
input_path=config.input_path.resolve(),
|
||||
database_path=database_path,
|
||||
output_path=output_path,
|
||||
episode_gap_days=config.episode_gap_days,
|
||||
memory_limit=config.memory_limit.upper(),
|
||||
threads=config.threads,
|
||||
allow_gaps=config.allow_gaps,
|
||||
allow_external_output=config.allow_external_output,
|
||||
overwrite=config.overwrite,
|
||||
)
|
||||
|
||||
|
||||
def _load_sql(name: str) -> str:
|
||||
path = SQL_ROOT / name
|
||||
try:
|
||||
return path.read_text(encoding="utf-8")
|
||||
except OSError as exc:
|
||||
raise MartBuildError(f"Cannot read local transform SQL: {path}") from exc
|
||||
|
||||
|
||||
def _remove_partial(paths: Iterable[Path]) -> None:
|
||||
for path in paths:
|
||||
path.unlink(missing_ok=True)
|
||||
|
||||
|
||||
def _output_columns(connection: duckdb.DuckDBPyConnection) -> List[Dict[str, str]]:
|
||||
rows = connection.execute("DESCRIBE SELECT * FROM feature_mart").fetchall()
|
||||
return [
|
||||
{"name": str(row[0]), "duckdb_type": str(row[1])}
|
||||
for row in rows
|
||||
]
|
||||
|
||||
|
||||
def build_feature_mart(config: BuildConfig) -> BuildSummary:
|
||||
config = _validate_config(config)
|
||||
batches = discover_and_validate_batches(config)
|
||||
source_data_kind = batches[0].export_kind
|
||||
population_estimate_allowed = (
|
||||
source_data_kind == BOUNDED_KIND and not config.allow_gaps
|
||||
)
|
||||
|
||||
database_path = config.database_path
|
||||
output_path = config.output_path
|
||||
manifest_path = output_path.with_name(output_path.name + ".manifest.json")
|
||||
partial_database = database_path.with_name(database_path.name + ".partial")
|
||||
partial_output = output_path.with_name(output_path.name + ".partial")
|
||||
partial_manifest = manifest_path.with_name(manifest_path.name + ".partial")
|
||||
partial_wal = partial_database.with_name(partial_database.name + ".wal")
|
||||
partial_paths = (partial_database, partial_output, partial_manifest, partial_wal)
|
||||
|
||||
finals = (database_path, output_path, manifest_path)
|
||||
existing = [path for path in finals if path.exists()]
|
||||
if existing and not config.overwrite:
|
||||
raise MartBuildError(
|
||||
"Refusing to overwrite existing outputs: "
|
||||
+ ", ".join(str(path) for path in existing)
|
||||
)
|
||||
|
||||
os.umask(0o077)
|
||||
database_path.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
|
||||
temp_directory = database_path.parent / ".duckdb_tmp"
|
||||
temp_directory.mkdir(parents=True, exist_ok=True, mode=0o700)
|
||||
_remove_partial(partial_paths)
|
||||
|
||||
connection: Optional[duckdb.DuckDBPyConnection] = None
|
||||
try:
|
||||
connection = duckdb.connect(str(partial_database))
|
||||
connection.execute("SET threads = ?", [config.threads])
|
||||
connection.execute("SET memory_limit = ?", [config.memory_limit])
|
||||
connection.execute("SET temp_directory = ?", [str(temp_directory)])
|
||||
connection.execute("SET preserve_insertion_order = false")
|
||||
connection.execute(
|
||||
"""
|
||||
CREATE TABLE build_config AS
|
||||
SELECT
|
||||
?::INTEGER AS episode_gap_days,
|
||||
?::VARCHAR AS source_data_kind,
|
||||
?::BOOLEAN AS population_estimate_allowed
|
||||
""",
|
||||
[
|
||||
config.episode_gap_days,
|
||||
source_data_kind,
|
||||
population_estimate_allowed,
|
||||
],
|
||||
)
|
||||
connection.execute(CREATE_STAGING_SQL)
|
||||
|
||||
for batch in batches:
|
||||
connection.execute(
|
||||
INGEST_SQL,
|
||||
[
|
||||
str(batch.path),
|
||||
batch.export_kind,
|
||||
batch.start.isoformat(sep=" "),
|
||||
batch.end.isoformat(sep=" "),
|
||||
str(batch.path),
|
||||
],
|
||||
)
|
||||
|
||||
actual_by_file = dict(
|
||||
connection.execute(
|
||||
"SELECT batch_file, count(*) FROM stg_events GROUP BY batch_file"
|
||||
).fetchall()
|
||||
)
|
||||
for batch in batches:
|
||||
actual = int(actual_by_file.get(str(batch.path), 0))
|
||||
if actual != batch.rows:
|
||||
raise MartBuildError(
|
||||
f"CSV row count mismatch for {batch.path}: "
|
||||
f"manifest={batch.rows}, actual={actual}"
|
||||
)
|
||||
|
||||
duplicate_ids = connection.execute(
|
||||
"""
|
||||
SELECT count(*)
|
||||
FROM (
|
||||
SELECT internal_event_id
|
||||
FROM stg_events
|
||||
GROUP BY internal_event_id
|
||||
HAVING count(*) > 1
|
||||
)
|
||||
"""
|
||||
).fetchone()[0]
|
||||
if duplicate_ids:
|
||||
raise MartBuildError(
|
||||
f"Input contains {duplicate_ids} duplicated internal event IDs"
|
||||
)
|
||||
|
||||
inconsistent_buckets = connection.execute(
|
||||
"""
|
||||
SELECT count(*)
|
||||
FROM (
|
||||
SELECT lower(vehicle_token) AS vehicle_token
|
||||
FROM stg_events
|
||||
WHERE regexp_full_match(lower(vehicle_token), '^[0-9a-f]{64}$')
|
||||
AND try_cast(vehicle_bucket AS INTEGER) BETWEEN 0 AND 99
|
||||
GROUP BY lower(vehicle_token)
|
||||
HAVING count(DISTINCT try_cast(vehicle_bucket AS INTEGER)) > 1
|
||||
)
|
||||
"""
|
||||
).fetchone()[0]
|
||||
if inconsistent_buckets:
|
||||
raise MartBuildError(
|
||||
f"Input contains {inconsistent_buckets} vehicle tokens with "
|
||||
"inconsistent audit buckets"
|
||||
)
|
||||
|
||||
sql_hashes: Dict[str, str] = {}
|
||||
for name in SQL_FILES:
|
||||
sql = _load_sql(name)
|
||||
sql_hashes[name] = hashlib.sha256(sql.encode("utf-8")).hexdigest()
|
||||
connection.execute(sql)
|
||||
|
||||
violations = connection.execute(
|
||||
"""
|
||||
SELECT check_name, violation_count
|
||||
FROM mart_validation
|
||||
WHERE violation_count <> 0
|
||||
ORDER BY check_name
|
||||
"""
|
||||
).fetchall()
|
||||
if violations:
|
||||
details = ", ".join(f"{name}={count}" for name, count in violations)
|
||||
raise MartBuildError(f"Mart invariant failure: {details}")
|
||||
|
||||
counts = connection.execute(
|
||||
"""
|
||||
SELECT
|
||||
(SELECT count(*) FROM stg_events),
|
||||
(SELECT count(*) FROM clean_events),
|
||||
(SELECT count(*) FROM inspection_episodes),
|
||||
(SELECT count(*) FROM feature_mart),
|
||||
(SELECT count(*) FROM feature_mart
|
||||
WHERE eligible_returning_target)
|
||||
"""
|
||||
).fetchone()
|
||||
input_rows, clean_events, episodes, mart_rows, eligible_rows = map(
|
||||
int, counts
|
||||
)
|
||||
|
||||
columns = _output_columns(connection)
|
||||
connection.execute("ANALYZE")
|
||||
connection.execute(
|
||||
"""
|
||||
COPY (
|
||||
SELECT *
|
||||
FROM feature_mart
|
||||
ORDER BY vehicle_token, episode_number
|
||||
) TO ? (FORMAT PARQUET, COMPRESSION ZSTD)
|
||||
""",
|
||||
[str(partial_output)],
|
||||
)
|
||||
|
||||
build_id = hashlib.sha256(
|
||||
json.dumps(
|
||||
{
|
||||
"inputs": [batch.compressed_sha256 for batch in batches],
|
||||
"sql": sql_hashes,
|
||||
"gap": config.episode_gap_days,
|
||||
},
|
||||
sort_keys=True,
|
||||
).encode("utf-8")
|
||||
).hexdigest()[:16]
|
||||
output_sha256 = file_sha256(partial_output)
|
||||
build_manifest: Dict[str, Any] = {
|
||||
"build_kind": "private_leakage_safe_inspection_feature_mart",
|
||||
"build_id": build_id,
|
||||
"generated_at_utc": datetime.now(timezone.utc).isoformat(),
|
||||
"classification": "private_pseudonymized_analytical_mart",
|
||||
"source_data_kind": source_data_kind,
|
||||
"population_estimate_allowed": population_estimate_allowed,
|
||||
"episode_gap_days": config.episode_gap_days,
|
||||
"vehicle_quality_contract": (
|
||||
"eligibility uses only prior events: at most 50 prior events "
|
||||
"and at most 4 prior events on any day"
|
||||
),
|
||||
"outcome_label_source_contract": {
|
||||
"allowed_values": ["overall_result", "utah_obd_proxy"],
|
||||
"null_allowed_only_for_unlabeled_outcomes": True,
|
||||
"utah_obd_proxy_source_era": "utah",
|
||||
"feature_role": "audit_only_not_a_predictor",
|
||||
},
|
||||
"vin_audit_contract": "vehicle_bucket < 10",
|
||||
"inputs": [
|
||||
{
|
||||
"file": str(batch.path),
|
||||
"manifest": str(batch.manifest_path),
|
||||
"export_kind": batch.export_kind,
|
||||
"query_version": batch.query_version,
|
||||
"query_sha256": batch.query_sha256,
|
||||
"source_start_inclusive": batch.start.isoformat(),
|
||||
"source_end_exclusive": batch.end.isoformat(),
|
||||
"rows": batch.rows,
|
||||
"compressed_file_sha256": batch.compressed_sha256,
|
||||
}
|
||||
for batch in batches
|
||||
],
|
||||
"vehicle_token_key_version": batches[0].key_version,
|
||||
"vehicle_token_key_fingerprint": batches[0].key_fingerprint,
|
||||
"transform_sql_sha256": sql_hashes,
|
||||
"row_counts": {
|
||||
"input": input_rows,
|
||||
"clean_events": clean_events,
|
||||
"episodes": episodes,
|
||||
"feature_mart": mart_rows,
|
||||
"eligible_returning_targets": eligible_rows,
|
||||
},
|
||||
"columns": columns,
|
||||
"parquet_sha256": output_sha256,
|
||||
"parquet_bytes": partial_output.stat().st_size,
|
||||
}
|
||||
partial_manifest.write_text(
|
||||
json.dumps(build_manifest, indent=2) + "\n", encoding="utf-8"
|
||||
)
|
||||
|
||||
connection.execute("CHECKPOINT")
|
||||
connection.close()
|
||||
connection = None
|
||||
|
||||
os.replace(partial_database, database_path)
|
||||
os.replace(partial_output, output_path)
|
||||
os.replace(partial_manifest, manifest_path)
|
||||
|
||||
return BuildSummary(
|
||||
database_path=database_path,
|
||||
output_path=output_path,
|
||||
manifest_path=manifest_path,
|
||||
input_rows=input_rows,
|
||||
clean_events=clean_events,
|
||||
episodes=episodes,
|
||||
mart_rows=mart_rows,
|
||||
eligible_rows=eligible_rows,
|
||||
source_data_kind=source_data_kind,
|
||||
population_estimate_allowed=population_estimate_allowed,
|
||||
)
|
||||
except Exception:
|
||||
if connection is not None:
|
||||
connection.close()
|
||||
_remove_partial(partial_paths)
|
||||
raise
|
||||
|
||||
|
||||
def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument(
|
||||
"--input",
|
||||
"--batch-dir",
|
||||
dest="input_path",
|
||||
type=Path,
|
||||
default=DEFAULT_INPUT,
|
||||
help="A bounded-batch directory or one development history CSV.gz.",
|
||||
)
|
||||
parser.add_argument("--database", type=Path, default=DEFAULT_DATABASE)
|
||||
parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
|
||||
parser.add_argument("--episode-gap-days", type=int, default=30)
|
||||
parser.add_argument("--memory-limit", default="4GB")
|
||||
parser.add_argument(
|
||||
"--threads", type=int, default=max(1, min(4, os.cpu_count() or 1))
|
||||
)
|
||||
parser.add_argument("--allow-gaps", action="store_true")
|
||||
parser.add_argument("--allow-external-output", action="store_true")
|
||||
parser.add_argument("--overwrite", action="store_true")
|
||||
return parser.parse_args(argv)
|
||||
|
||||
|
||||
def main(argv: Optional[Sequence[str]] = None) -> int:
|
||||
args = parse_args(argv)
|
||||
config = BuildConfig(
|
||||
input_path=args.input_path,
|
||||
database_path=args.database,
|
||||
output_path=args.output,
|
||||
episode_gap_days=args.episode_gap_days,
|
||||
memory_limit=args.memory_limit,
|
||||
threads=args.threads,
|
||||
allow_gaps=args.allow_gaps,
|
||||
allow_external_output=args.allow_external_output,
|
||||
overwrite=args.overwrite,
|
||||
)
|
||||
try:
|
||||
summary = build_feature_mart(config)
|
||||
except (MartBuildError, duckdb.Error, OSError) as exc:
|
||||
print(f"Feature-mart build failed: {exc}", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
print(
|
||||
f"Built {summary.mart_rows:,} episode rows "
|
||||
f"({summary.eligible_rows:,} eligible returning targets)"
|
||||
)
|
||||
print(f"Private DuckDB warehouse: {summary.database_path}")
|
||||
print(f"Private Parquet mart: {summary.output_path}")
|
||||
print(f"Build manifest: {summary.manifest_path}")
|
||||
if not summary.population_estimate_allowed:
|
||||
print("Development sample: population estimates are prohibited.")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,96 @@
|
||||
"""Verify the countydata connection without exposing credentials or row data."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
import psycopg2
|
||||
|
||||
|
||||
REQUIRED_ENV_VARS = (
|
||||
"PGHOST",
|
||||
"PGPORT",
|
||||
"PGDATABASE",
|
||||
"PGUSER",
|
||||
"PGPASSWORD",
|
||||
"PGSSLMODE",
|
||||
)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
missing = [name for name in REQUIRED_ENV_VARS if not os.environ.get(name)]
|
||||
if missing:
|
||||
print(
|
||||
"Missing database environment variables: " + ", ".join(missing),
|
||||
file=sys.stderr,
|
||||
)
|
||||
print("Run this through the VS Code task so .env is loaded safely.", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
if os.environ["PGSSLMODE"].lower() not in {"require", "verify-ca", "verify-full"}:
|
||||
print("Refusing to connect because PGSSLMODE does not require TLS.", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
try:
|
||||
connection = psycopg2.connect(
|
||||
application_name="summer_project_vscode_check",
|
||||
connect_timeout=10,
|
||||
options=(
|
||||
"-c default_transaction_read_only=on "
|
||||
"-c statement_timeout=30000"
|
||||
),
|
||||
)
|
||||
connection.set_session(readonly=True)
|
||||
except psycopg2.Error as exc:
|
||||
print(f"Connection failed: {exc.diag.message_primary or type(exc).__name__}", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
try:
|
||||
with connection, connection.cursor() as cursor:
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT
|
||||
current_database(),
|
||||
current_user,
|
||||
current_setting('transaction_read_only'),
|
||||
version()
|
||||
"""
|
||||
)
|
||||
database, user, read_only, version = cursor.fetchone()
|
||||
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT ssl, version, cipher, bits
|
||||
FROM pg_stat_ssl
|
||||
WHERE pid = pg_backend_pid()
|
||||
"""
|
||||
)
|
||||
ssl, tls_version, cipher, bits = cursor.fetchone()
|
||||
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT count(*)
|
||||
FROM pg_class AS c
|
||||
JOIN pg_namespace AS n ON n.oid = c.relnamespace
|
||||
WHERE c.relkind IN ('r', 'p', 'v', 'm', 'f')
|
||||
AND n.nspname NOT IN ('pg_catalog', 'information_schema')
|
||||
AND n.nspname NOT LIKE 'pg_toast%'
|
||||
"""
|
||||
)
|
||||
relation_count = cursor.fetchone()[0]
|
||||
|
||||
print("Countydata connection succeeded")
|
||||
print(f" database: {database}")
|
||||
print(f" user: {user}")
|
||||
print(f" transaction read-only: {read_only}")
|
||||
print(f" TLS: {ssl} ({tls_version}, {cipher}, {bits}-bit)")
|
||||
print(f" visible non-system relations: {relation_count}")
|
||||
print(f" server: {version}")
|
||||
return 0
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,56 @@
|
||||
"""Create the local project-specific VIN HMAC key without printing it."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import os
|
||||
import secrets
|
||||
import stat
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
||||
DEFAULT_PATH = PROJECT_ROOT / ".secrets/vin_hmac.key"
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--path", type=Path, default=DEFAULT_PATH)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def key_fingerprint(key: bytes) -> str:
|
||||
return hashlib.sha256(
|
||||
b"utah-vehicle-health/key-fingerprint/v1\0" + key
|
||||
).hexdigest()[:16]
|
||||
|
||||
|
||||
def main() -> int:
|
||||
path = parse_args().path.resolve()
|
||||
os.umask(0o077)
|
||||
path.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
|
||||
|
||||
if path.exists():
|
||||
mode = stat.S_IMODE(path.stat().st_mode)
|
||||
if mode & 0o077:
|
||||
print(f"Existing key has unsafe permissions: {path}", file=sys.stderr)
|
||||
return 1
|
||||
print(f"Key already exists; leaving it unchanged: {path}")
|
||||
return 0
|
||||
|
||||
key = secrets.token_bytes(32)
|
||||
descriptor = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
|
||||
with os.fdopen(descriptor, "w", encoding="ascii") as handle:
|
||||
handle.write(key.hex() + "\n")
|
||||
handle.flush()
|
||||
os.fsync(handle.fileno())
|
||||
|
||||
print(f"Created private VIN HMAC key: {path}")
|
||||
print(f"Non-secret key fingerprint: {key_fingerprint(key)}")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,234 @@
|
||||
"""Export complete histories for a private development-only vehicle sample.
|
||||
|
||||
The source sample is page-based and therefore not population-representative.
|
||||
Use it to develop and test the episode/model pipeline, never for final rates.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import gzip
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import psycopg2
|
||||
|
||||
try:
|
||||
from . import export_inspection_batch as secure_export
|
||||
except ImportError: # Direct execution: python scripts/export_history_sample.py
|
||||
import export_inspection_batch as secure_export
|
||||
|
||||
|
||||
QUERY_VERSION = "inspection_history_development_sample_v3"
|
||||
|
||||
|
||||
def parse_timestamp(value: str) -> datetime:
|
||||
try:
|
||||
return datetime.fromisoformat(value)
|
||||
except ValueError as exc:
|
||||
raise argparse.ArgumentTypeError(f"Invalid ISO timestamp: {value}") from exc
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--vehicles", type=int, default=10_000)
|
||||
parser.add_argument("--sample-percent", type=float, default=0.5)
|
||||
parser.add_argument("--seed", type=int, default=20260715)
|
||||
parser.add_argument("--start", type=parse_timestamp, default=datetime(2010, 1, 1))
|
||||
parser.add_argument("--end", type=parse_timestamp, default=datetime(2026, 7, 1))
|
||||
parser.add_argument("--fetch-size", type=int, default=5_000)
|
||||
parser.add_argument("--output", type=Path)
|
||||
parser.add_argument("--overwrite", action="store_true")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def validate_args(args: argparse.Namespace) -> Path:
|
||||
if not 100 <= args.vehicles <= 100_000:
|
||||
raise ValueError("--vehicles must be between 100 and 100000")
|
||||
if not 0 < args.sample_percent <= 5:
|
||||
raise ValueError("--sample-percent must be greater than 0 and at most 5")
|
||||
if not 100 <= args.fetch_size <= 50_000:
|
||||
raise ValueError("--fetch-size must be between 100 and 50000")
|
||||
if args.end <= args.start:
|
||||
raise ValueError("--end must be later than --start")
|
||||
|
||||
output = args.output
|
||||
if output is None:
|
||||
output = (
|
||||
secure_export.PRIVATE_DATA_ROOT
|
||||
/ "development"
|
||||
/ f"history_sample_{args.vehicles}.csv.gz"
|
||||
)
|
||||
output = output.resolve()
|
||||
if not output.is_relative_to(secure_export.PRIVATE_DATA_ROOT):
|
||||
raise ValueError(
|
||||
f"Development histories must stay under {secure_export.PRIVATE_DATA_ROOT}"
|
||||
)
|
||||
manifest = output.with_name(output.name + ".manifest.json")
|
||||
if (output.exists() or manifest.exists()) and not args.overwrite:
|
||||
raise FileExistsError(
|
||||
f"Refusing to overwrite {output} or its manifest; use --overwrite"
|
||||
)
|
||||
return output
|
||||
|
||||
|
||||
def make_sql(sample_percent: float, seed: int) -> str:
|
||||
safe_percent = f"{sample_percent:.6f}"
|
||||
safe_seed = int(seed)
|
||||
return f"""
|
||||
WITH candidate_vins AS MATERIALIZED (
|
||||
SELECT DISTINCT vin
|
||||
FROM data.inspection_search
|
||||
TABLESAMPLE SYSTEM ({safe_percent}) REPEATABLE ({safe_seed})
|
||||
WHERE vin IS NOT NULL
|
||||
AND length(btrim(vin)) = 17
|
||||
),
|
||||
sampled_vins AS MATERIALIZED (
|
||||
SELECT vin
|
||||
FROM candidate_vins
|
||||
ORDER BY md5(vin)
|
||||
LIMIT %(vehicle_limit)s
|
||||
)
|
||||
SELECT
|
||||
s.id AS internal_event_id,
|
||||
s.vin,
|
||||
s.test_start AS event_ts,
|
||||
lower(btrim(s.county)) AS source_era,
|
||||
CASE
|
||||
WHEN lower(btrim(s.county)) IN ('slc', 'slco') THEN 'salt_lake'
|
||||
ELSE lower(btrim(s.county))
|
||||
END AS public_county,
|
||||
{secure_export.OUTCOME_SELECT_SQL},
|
||||
lower(nullif(btrim(s.program_type), '')) AS program_type,
|
||||
upper(nullif(btrim(s.test_type), '')) AS test_type,
|
||||
upper(nullif(btrim(v.make), '')) AS observed_make,
|
||||
upper(nullif(btrim(v.model), '')) AS observed_model,
|
||||
v.year AS observed_model_year
|
||||
FROM data.inspection_search AS s
|
||||
JOIN data.inspection_vehicle AS v USING (id)
|
||||
JOIN sampled_vins ON sampled_vins.vin = s.vin
|
||||
WHERE s.test_start >= %(start_ts)s
|
||||
AND s.test_start < %(end_ts)s
|
||||
ORDER BY s.test_start, s.id
|
||||
"""
|
||||
|
||||
|
||||
def file_sha256(path: Path) -> str:
|
||||
digest = hashlib.sha256()
|
||||
with path.open("rb") as handle:
|
||||
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
|
||||
digest.update(chunk)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def main() -> int:
|
||||
args = parse_args()
|
||||
os.umask(0o077)
|
||||
try:
|
||||
output = validate_args(args)
|
||||
key, key_version, key_fingerprint = secure_export.get_hmac_configuration()
|
||||
except (FileExistsError, ValueError) as exc:
|
||||
print(f"Configuration error: {exc}", file=sys.stderr)
|
||||
return 2
|
||||
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
manifest = output.with_name(output.name + ".manifest.json")
|
||||
partial = output.with_name(output.name + ".partial")
|
||||
partial_manifest = manifest.with_name(manifest.name + ".partial")
|
||||
sql = make_sql(args.sample_percent, args.seed)
|
||||
|
||||
try:
|
||||
connection = secure_export.connect_read_only()
|
||||
except (ValueError, RuntimeError, psycopg2.Error) as exc:
|
||||
print(f"Secure read-only connection failed: {type(exc).__name__}", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
source_rows_read = 0
|
||||
rows_written = 0
|
||||
invalid_vins_skipped = 0
|
||||
last_order_key: Optional[tuple[object, object]] = None
|
||||
try:
|
||||
with gzip.open(partial, "wt", encoding="utf-8", newline="") as handle:
|
||||
writer = csv.writer(handle)
|
||||
writer.writerow(secure_export.OUTPUT_COLUMNS)
|
||||
with connection.cursor(name="uvh_history_sample") as cursor:
|
||||
cursor.itersize = args.fetch_size
|
||||
cursor.execute(
|
||||
sql,
|
||||
{
|
||||
"vehicle_limit": args.vehicles,
|
||||
"start_ts": args.start,
|
||||
"end_ts": args.end,
|
||||
},
|
||||
)
|
||||
while True:
|
||||
rows = cursor.fetchmany(args.fetch_size)
|
||||
if not rows:
|
||||
break
|
||||
for source_row in rows:
|
||||
source_rows_read += 1
|
||||
order_key = (source_row[2], source_row[0])
|
||||
if last_order_key is not None and order_key <= last_order_key:
|
||||
raise RuntimeError(
|
||||
"Source join did not return a unique increasing "
|
||||
"event_ts/internal_event_id order"
|
||||
)
|
||||
last_order_key = order_key
|
||||
transformed = secure_export.private_row(
|
||||
source_row, key, key_version
|
||||
)
|
||||
if transformed is None:
|
||||
invalid_vins_skipped += 1
|
||||
continue
|
||||
writer.writerow(transformed)
|
||||
rows_written += 1
|
||||
connection.rollback()
|
||||
except Exception:
|
||||
partial.unlink(missing_ok=True)
|
||||
partial_manifest.unlink(missing_ok=True)
|
||||
raise
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
manifest_data = {
|
||||
"export_kind": "vehicle_history_development_sample",
|
||||
"query_version": QUERY_VERSION,
|
||||
"query_sha256": hashlib.sha256(sql.encode("utf-8")).hexdigest(),
|
||||
"generated_at_utc": datetime.now(timezone.utc).isoformat(),
|
||||
"source_start_inclusive": args.start.isoformat(),
|
||||
"source_end_exclusive": args.end.isoformat(),
|
||||
"vehicle_sample_limit": args.vehicles,
|
||||
"sample_method": "page_sample_candidates_then_md5_ordered_distinct_vins",
|
||||
"sample_percent": args.sample_percent,
|
||||
"sample_seed": args.seed,
|
||||
"population_estimate_allowed": False,
|
||||
"vehicle_token_key_version": key_version,
|
||||
"vehicle_token_key_fingerprint": key_fingerprint,
|
||||
"source_rows_read": source_rows_read,
|
||||
"rows": rows_written,
|
||||
"invalid_vins_skipped": invalid_vins_skipped,
|
||||
"columns": secure_export.OUTPUT_COLUMNS,
|
||||
"compressed_file_sha256": file_sha256(partial),
|
||||
"compressed_file_bytes": partial.stat().st_size,
|
||||
"classification": "private_pseudonymized_development_only",
|
||||
"source_transaction": "read_only_repeatable_read",
|
||||
}
|
||||
partial_manifest.write_text(
|
||||
json.dumps(manifest_data, indent=2) + "\n", encoding="utf-8"
|
||||
)
|
||||
os.replace(partial, output)
|
||||
os.replace(partial_manifest, manifest)
|
||||
print(f"Exported {rows_written:,} rows for a development vehicle sample")
|
||||
print(f"Private output: {output}")
|
||||
print("This sample is not valid for population-rate estimates.")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,414 @@
|
||||
"""Export one private inspection batch while replacing VINs with keyed HMACs.
|
||||
|
||||
The output is private analytical staging data, not a public-dashboard dataset.
|
||||
It intentionally omits plates, ZIPs, stations, raw JSON, and operational data.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import gzip
|
||||
import hashlib
|
||||
import hmac
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import stat
|
||||
import sys
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Iterable, Optional, Sequence
|
||||
|
||||
import psycopg2
|
||||
|
||||
|
||||
QUERY_VERSION = "inspection_batch_v4"
|
||||
MAX_BATCH_DAYS = 32
|
||||
SECURE_SSL_MODES = {"require", "verify-ca", "verify-full"}
|
||||
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
||||
DEFAULT_KEY_FILE = PROJECT_ROOT / ".secrets/vin_hmac.key"
|
||||
PRIVATE_DATA_ROOT = (PROJECT_ROOT / "data/private").resolve()
|
||||
VIN_PATTERN = re.compile(r"^[A-HJ-NPR-Z0-9]{17}$")
|
||||
KEY_VERSION_PATTERN = re.compile(r"^[A-Za-z0-9._-]{1,32}$")
|
||||
|
||||
OUTCOME_ALIASES = {
|
||||
"PASS": "pass",
|
||||
"P": "pass",
|
||||
"FAIL": "fail",
|
||||
"F": "fail",
|
||||
"REJECT": "reject",
|
||||
"ABORT": "abort",
|
||||
}
|
||||
|
||||
|
||||
def _normalized_outcome(value: object) -> Optional[str]:
|
||||
if not isinstance(value, str):
|
||||
return None
|
||||
return OUTCOME_ALIASES.get(value.strip().upper())
|
||||
|
||||
|
||||
def canonicalize_outcome(
|
||||
overall_result: object,
|
||||
obd_result: object,
|
||||
source_era: object,
|
||||
program_type: object = None,
|
||||
test_type: object = None,
|
||||
) -> tuple[Optional[str], Optional[str]]:
|
||||
"""Return the approved label and its provenance for one inspection."""
|
||||
overall = _normalized_outcome(overall_result)
|
||||
if overall is not None:
|
||||
return overall, "overall_result"
|
||||
|
||||
proxy_eligible = (
|
||||
isinstance(source_era, str)
|
||||
and source_era.strip().lower() == "utah"
|
||||
and isinstance(program_type, str)
|
||||
and program_type.strip().lower() == "obd"
|
||||
and isinstance(test_type, str)
|
||||
and test_type.strip().upper() == "OBD"
|
||||
)
|
||||
if proxy_eligible:
|
||||
proxy = _normalized_outcome(obd_result)
|
||||
if proxy is not None:
|
||||
# This narrowly validated Utah OBD/OBD proxy is approved only for
|
||||
# binary pass/non-pass labels. Multiclass work must retain
|
||||
# overall_result labels exclusively.
|
||||
return proxy, "utah_obd_proxy"
|
||||
return None, None
|
||||
|
||||
|
||||
def _outcome_case_sql(column: str) -> str:
|
||||
cases = " ".join(
|
||||
f"WHEN '{raw}' THEN '{canonical}'"
|
||||
for raw, canonical in OUTCOME_ALIASES.items()
|
||||
)
|
||||
return f"CASE upper(btrim({column})) {cases} ELSE NULL END"
|
||||
|
||||
|
||||
_OVERALL_OUTCOME_SQL = _outcome_case_sql("s.overall_result")
|
||||
_UTAH_OBD_OUTCOME_SQL = _outcome_case_sql("s.obd_result")
|
||||
_UTAH_OBD_PROXY_PREDICATE_SQL = """lower(btrim(s.county)) = 'utah'
|
||||
AND lower(btrim(s.program_type)) = 'obd'
|
||||
AND upper(btrim(s.test_type)) = 'OBD'"""
|
||||
|
||||
# Keep label provenance in private staging. Only Utah rows normalized to both
|
||||
# program_type=obd and test_type=OBD may use the proxy, and then only for binary
|
||||
# pass/non-pass modeling. Multiclass labels must use overall_result.
|
||||
OUTCOME_SELECT_SQL = f"""CASE
|
||||
WHEN {_OVERALL_OUTCOME_SQL} IS NOT NULL THEN {_OVERALL_OUTCOME_SQL}
|
||||
WHEN {_UTAH_OBD_PROXY_PREDICATE_SQL} THEN {_UTAH_OBD_OUTCOME_SQL}
|
||||
ELSE NULL
|
||||
END AS canonical_outcome,
|
||||
CASE
|
||||
WHEN {_OVERALL_OUTCOME_SQL} IS NOT NULL THEN 'overall_result'
|
||||
WHEN {_UTAH_OBD_PROXY_PREDICATE_SQL}
|
||||
AND {_UTAH_OBD_OUTCOME_SQL} IS NOT NULL THEN 'utah_obd_proxy'
|
||||
ELSE NULL
|
||||
END AS outcome_label_source"""
|
||||
|
||||
|
||||
EXTRACT_SQL = f"""
|
||||
SELECT
|
||||
s.id AS internal_event_id,
|
||||
s.vin,
|
||||
s.test_start AS event_ts,
|
||||
lower(btrim(s.county)) AS source_era,
|
||||
CASE
|
||||
WHEN lower(btrim(s.county)) IN ('slc', 'slco') THEN 'salt_lake'
|
||||
ELSE lower(btrim(s.county))
|
||||
END AS public_county,
|
||||
{OUTCOME_SELECT_SQL},
|
||||
lower(nullif(btrim(s.program_type), '')) AS program_type,
|
||||
upper(nullif(btrim(s.test_type), '')) AS test_type,
|
||||
upper(nullif(btrim(v.make), '')) AS observed_make,
|
||||
upper(nullif(btrim(v.model), '')) AS observed_model,
|
||||
v.year AS observed_model_year
|
||||
FROM data.inspection_search AS s
|
||||
JOIN data.inspection_vehicle AS v USING (id)
|
||||
WHERE s.test_start >= %(start_ts)s
|
||||
AND s.test_start < %(end_ts)s
|
||||
AND s.vin IS NOT NULL
|
||||
AND length(btrim(s.vin)) = 17
|
||||
ORDER BY s.test_start, s.id
|
||||
"""
|
||||
|
||||
OUTPUT_COLUMNS = (
|
||||
"internal_event_id",
|
||||
"vehicle_token",
|
||||
"vehicle_bucket",
|
||||
"event_ts",
|
||||
"source_era",
|
||||
"public_county",
|
||||
"canonical_outcome",
|
||||
"outcome_label_source",
|
||||
"program_type",
|
||||
"test_type",
|
||||
"observed_make",
|
||||
"observed_model",
|
||||
"observed_model_year",
|
||||
)
|
||||
|
||||
|
||||
def parse_timestamp(value: str) -> datetime:
|
||||
try:
|
||||
return datetime.fromisoformat(value)
|
||||
except ValueError as exc:
|
||||
raise argparse.ArgumentTypeError(
|
||||
f"Invalid ISO timestamp or date: {value}"
|
||||
) from exc
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--start", required=True, type=parse_timestamp)
|
||||
parser.add_argument("--end", required=True, type=parse_timestamp)
|
||||
parser.add_argument("--output", type=Path)
|
||||
parser.add_argument("--page-size", type=int, default=10_000)
|
||||
parser.add_argument("--overwrite", action="store_true")
|
||||
parser.add_argument(
|
||||
"--allow-external-output",
|
||||
action="store_true",
|
||||
help="Allow private row-level output outside data/private (unsafe).",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def validate_args(args: argparse.Namespace) -> Path:
|
||||
if args.end <= args.start:
|
||||
raise ValueError("--end must be later than --start")
|
||||
if (args.end - args.start).total_seconds() > MAX_BATCH_DAYS * 86400:
|
||||
raise ValueError(
|
||||
f"A source batch may span at most {MAX_BATCH_DAYS} days; "
|
||||
"export longer periods as closed monthly batches."
|
||||
)
|
||||
if not 100 <= args.page_size <= 50_000:
|
||||
raise ValueError("--page-size must be between 100 and 50000")
|
||||
|
||||
output = args.output
|
||||
if output is None:
|
||||
name = f"inspection_{args.start:%Y%m%d}_{args.end:%Y%m%d}.csv.gz"
|
||||
output = PRIVATE_DATA_ROOT / "inspection_batches" / name
|
||||
output = output.resolve()
|
||||
if not args.allow_external_output and not output.is_relative_to(PRIVATE_DATA_ROOT):
|
||||
raise ValueError(
|
||||
f"Private exports must stay under {PRIVATE_DATA_ROOT}; "
|
||||
"use --allow-external-output only for controlled validation."
|
||||
)
|
||||
manifest = output.with_name(output.name + ".manifest.json")
|
||||
if (output.exists() or manifest.exists()) and not args.overwrite:
|
||||
raise FileExistsError(
|
||||
f"Refusing to overwrite {output} or its manifest; use --overwrite"
|
||||
)
|
||||
return output
|
||||
|
||||
|
||||
def get_hmac_configuration() -> tuple[bytes, str, str]:
|
||||
encoded_key = os.environ.get("VIN_HASH_KEY", "").strip()
|
||||
key_file = Path(os.environ.get("VIN_HASH_KEY_FILE", DEFAULT_KEY_FILE))
|
||||
if not encoded_key and key_file.exists():
|
||||
mode = stat.S_IMODE(key_file.stat().st_mode)
|
||||
if mode & 0o077:
|
||||
raise ValueError(f"VIN HMAC key file must be mode 0600: {key_file}")
|
||||
encoded_key = key_file.read_text(encoding="ascii").strip()
|
||||
|
||||
version = os.environ.get("VIN_HASH_KEY_VERSION", "v1").strip() or "v1"
|
||||
if not KEY_VERSION_PATTERN.fullmatch(version):
|
||||
raise ValueError("VIN_HASH_KEY_VERSION contains unsupported characters")
|
||||
try:
|
||||
key = bytes.fromhex(encoded_key)
|
||||
except ValueError as exc:
|
||||
raise ValueError("VIN_HASH_KEY must be a hexadecimal value") from exc
|
||||
if len(key) != 32:
|
||||
raise ValueError(
|
||||
"VIN_HASH_KEY must encode exactly 32 random bytes (64 hex characters). "
|
||||
"See .env.example."
|
||||
)
|
||||
if len(set(key)) < 8:
|
||||
raise ValueError(
|
||||
"VIN_HASH_KEY has insufficient byte diversity; generate 32 random bytes."
|
||||
)
|
||||
fingerprint = hashlib.sha256(
|
||||
b"utah-vehicle-health/key-fingerprint/v1\0" + key
|
||||
).hexdigest()[:16]
|
||||
return key, version, fingerprint
|
||||
|
||||
|
||||
def normalize_vin(vin: str) -> Optional[str]:
|
||||
normalized = vin.strip().upper()
|
||||
if not VIN_PATTERN.fullmatch(normalized):
|
||||
return None
|
||||
if len(set(normalized)) < 4:
|
||||
return None
|
||||
if normalized in {"12345678901234567", "98765432109876543"}:
|
||||
return None
|
||||
return normalized
|
||||
|
||||
|
||||
def vehicle_token(vin: str, key: bytes, key_version: str) -> str:
|
||||
message = (
|
||||
f"utah-vehicle-health/{key_version}/vin\0".encode("ascii")
|
||||
+ vin.encode("ascii")
|
||||
)
|
||||
return hmac.new(key, message, hashlib.sha256).hexdigest()
|
||||
|
||||
|
||||
def private_row(
|
||||
row: Sequence[object], key: bytes, key_version: str
|
||||
) -> Optional[tuple[object, ...]]:
|
||||
internal_event_id, vin, *remaining = row
|
||||
if not isinstance(vin, str):
|
||||
raise ValueError("The source returned a non-text VIN")
|
||||
normalized = normalize_vin(vin)
|
||||
if normalized is None:
|
||||
return None
|
||||
token = vehicle_token(normalized, key, key_version)
|
||||
bucket = int(token[:8], 16) % 100
|
||||
return (internal_event_id, token, bucket, *remaining)
|
||||
|
||||
|
||||
def batches(
|
||||
connection: "psycopg2.extensions.connection",
|
||||
start: datetime,
|
||||
end: datetime,
|
||||
page_size: int,
|
||||
) -> Iterable[list[tuple[object, ...]]]:
|
||||
with connection.cursor(name="uvh_inspection_batch") as cursor:
|
||||
cursor.itersize = page_size
|
||||
cursor.execute(EXTRACT_SQL, {"start_ts": start, "end_ts": end})
|
||||
while True:
|
||||
rows = cursor.fetchmany(page_size)
|
||||
if not rows:
|
||||
return
|
||||
yield rows
|
||||
|
||||
|
||||
def connect_read_only() -> "psycopg2.extensions.connection":
|
||||
ssl_mode = os.environ.get("PGSSLMODE", "").lower()
|
||||
if ssl_mode not in SECURE_SSL_MODES:
|
||||
raise ValueError("PGSSLMODE must require TLS before exporting data")
|
||||
|
||||
connection = psycopg2.connect(
|
||||
application_name="utah_vehicle_health_batch_export",
|
||||
connect_timeout=10,
|
||||
options=(
|
||||
"-c default_transaction_read_only=on "
|
||||
"-c statement_timeout=120000 "
|
||||
"-c lock_timeout=3000 "
|
||||
"-c idle_in_transaction_session_timeout=60000"
|
||||
),
|
||||
)
|
||||
connection.set_session(
|
||||
readonly=True,
|
||||
isolation_level="REPEATABLE READ",
|
||||
autocommit=False,
|
||||
)
|
||||
with connection.cursor() as cursor:
|
||||
cursor.execute(
|
||||
"""
|
||||
SELECT
|
||||
current_setting('transaction_read_only'),
|
||||
ssl
|
||||
FROM pg_stat_ssl
|
||||
WHERE pid = pg_backend_pid()
|
||||
"""
|
||||
)
|
||||
read_only, ssl = cursor.fetchone()
|
||||
if read_only != "on" or not ssl:
|
||||
connection.close()
|
||||
raise RuntimeError("The database session is not read-only over TLS")
|
||||
return connection
|
||||
|
||||
|
||||
def main() -> int:
|
||||
args = parse_args()
|
||||
os.umask(0o077)
|
||||
try:
|
||||
output = validate_args(args)
|
||||
key, key_version, key_fingerprint = get_hmac_configuration()
|
||||
except (FileExistsError, ValueError) as exc:
|
||||
print(f"Configuration error: {exc}", file=sys.stderr)
|
||||
return 2
|
||||
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
partial = output.with_name(output.name + ".partial")
|
||||
manifest = output.with_name(output.name + ".manifest.json")
|
||||
partial_manifest = manifest.with_name(manifest.name + ".partial")
|
||||
|
||||
try:
|
||||
connection = connect_read_only()
|
||||
except (ValueError, RuntimeError, psycopg2.Error) as exc:
|
||||
print(f"Secure read-only connection failed: {type(exc).__name__}", file=sys.stderr)
|
||||
return 1
|
||||
|
||||
source_rows_read = 0
|
||||
rows_written = 0
|
||||
invalid_vins_skipped = 0
|
||||
last_order_key: Optional[tuple[object, object]] = None
|
||||
try:
|
||||
with gzip.open(partial, "wt", encoding="utf-8", newline="") as handle:
|
||||
writer = csv.writer(handle)
|
||||
writer.writerow(OUTPUT_COLUMNS)
|
||||
for source_rows in batches(
|
||||
connection, args.start, args.end, args.page_size
|
||||
):
|
||||
for source_row in source_rows:
|
||||
source_rows_read += 1
|
||||
order_key = (source_row[2], source_row[0])
|
||||
if last_order_key is not None and order_key <= last_order_key:
|
||||
raise RuntimeError(
|
||||
"Source join did not return a unique increasing "
|
||||
"event_ts/internal_event_id order"
|
||||
)
|
||||
last_order_key = order_key
|
||||
transformed = private_row(source_row, key, key_version)
|
||||
if transformed is None:
|
||||
invalid_vins_skipped += 1
|
||||
continue
|
||||
writer.writerow(transformed)
|
||||
rows_written += 1
|
||||
connection.rollback()
|
||||
except Exception:
|
||||
partial.unlink(missing_ok=True)
|
||||
partial_manifest.unlink(missing_ok=True)
|
||||
raise
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
digest = hashlib.sha256()
|
||||
with partial.open("rb") as handle:
|
||||
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
|
||||
digest.update(chunk)
|
||||
|
||||
manifest_data = {
|
||||
"export_kind": "bounded_batch",
|
||||
"query_version": QUERY_VERSION,
|
||||
"query_sha256": hashlib.sha256(EXTRACT_SQL.encode("utf-8")).hexdigest(),
|
||||
"generated_at_utc": datetime.now(timezone.utc).isoformat(),
|
||||
"source_start_inclusive": args.start.isoformat(),
|
||||
"source_end_exclusive": args.end.isoformat(),
|
||||
"vehicle_token_key_version": key_version,
|
||||
"vehicle_token_key_fingerprint": key_fingerprint,
|
||||
"source_rows_read": source_rows_read,
|
||||
"rows": rows_written,
|
||||
"invalid_vins_skipped": invalid_vins_skipped,
|
||||
"columns": OUTPUT_COLUMNS,
|
||||
"compressed_file_sha256": digest.hexdigest(),
|
||||
"compressed_file_bytes": partial.stat().st_size,
|
||||
"classification": "private_pseudonymized_analytical_staging",
|
||||
"source_transaction": "read_only_repeatable_read",
|
||||
}
|
||||
partial_manifest.write_text(
|
||||
json.dumps(manifest_data, indent=2) + "\n", encoding="utf-8"
|
||||
)
|
||||
os.replace(partial, output)
|
||||
os.replace(partial_manifest, manifest)
|
||||
print(f"Exported {rows_written:,} private pseudonymized rows to {output}")
|
||||
if invalid_vins_skipped:
|
||||
print(f"Skipped {invalid_vins_skipped:,} invalid/placeholder VIN rows")
|
||||
print(f"Wrote manifest to {manifest}")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,71 @@
|
||||
"""Run the project's fixed aggregate-only feasibility query safely.
|
||||
|
||||
This runner accepts no arbitrary SQL. It executes the reviewed query inside the
|
||||
same TLS-protected, read-only transaction used by the private exporters and
|
||||
prints only its aggregate result table.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import psycopg2
|
||||
|
||||
import export_inspection_batch as secure_export
|
||||
|
||||
|
||||
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
||||
QUERY_PATH = PROJECT_ROOT / "sql/10_episode_cohort_feasibility.sql"
|
||||
|
||||
|
||||
def reviewed_query() -> str:
|
||||
sql = QUERY_PATH.read_text(encoding="utf-8")
|
||||
begin = "BEGIN TRANSACTION READ ONLY;"
|
||||
rollback = "ROLLBACK;"
|
||||
if sql.count(begin) != 1 or sql.count(rollback) != 1:
|
||||
raise RuntimeError("The reviewed feasibility query wrappers changed")
|
||||
return sql.replace(begin, "", 1).rsplit(rollback, 1)[0].strip()
|
||||
|
||||
|
||||
def main() -> int:
|
||||
try:
|
||||
connection = secure_export.connect_read_only()
|
||||
except (ValueError, RuntimeError, psycopg2.Error) as exc:
|
||||
print(
|
||||
f"Secure read-only connection failed: {type(exc).__name__}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return 1
|
||||
|
||||
try:
|
||||
with connection.cursor() as cursor:
|
||||
cursor.execute(reviewed_query())
|
||||
columns = [description.name for description in cursor.description]
|
||||
rows = cursor.fetchall()
|
||||
connection.rollback()
|
||||
except psycopg2.Error as exc:
|
||||
connection.rollback()
|
||||
print(
|
||||
f"Aggregate feasibility query failed: "
|
||||
f"{exc.diag.message_primary or type(exc).__name__}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return 1
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
widths = [len(column) for column in columns]
|
||||
rendered = [["" if value is None else str(value) for value in row] for row in rows]
|
||||
for row in rendered:
|
||||
widths = [max(width, len(value)) for width, value in zip(widths, row)]
|
||||
|
||||
print(" ".join(value.ljust(width) for value, width in zip(columns, widths)))
|
||||
print(" ".join("-" * width for width in widths))
|
||||
for row in rendered:
|
||||
print(" ".join(value.ljust(width) for value, width in zip(row, widths)))
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,65 @@
|
||||
"""Run the fixed aggregate-only source outcome-mapping audit."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import psycopg2
|
||||
|
||||
import export_inspection_batch as secure_export
|
||||
|
||||
|
||||
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
||||
QUERY_PATH = PROJECT_ROOT / "sql/11_outcome_mapping_audit.sql"
|
||||
|
||||
|
||||
def reviewed_query() -> str:
|
||||
sql = QUERY_PATH.read_text(encoding="utf-8")
|
||||
begin = "BEGIN TRANSACTION READ ONLY;"
|
||||
rollback = "ROLLBACK;"
|
||||
if sql.count(begin) != 1 or sql.count(rollback) != 1:
|
||||
raise RuntimeError("The reviewed outcome audit wrappers changed")
|
||||
return sql.replace(begin, "", 1).rsplit(rollback, 1)[0].strip()
|
||||
|
||||
|
||||
def main() -> int:
|
||||
try:
|
||||
connection = secure_export.connect_read_only()
|
||||
except (ValueError, RuntimeError, psycopg2.Error) as exc:
|
||||
print(
|
||||
f"Secure read-only connection failed: {type(exc).__name__}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return 1
|
||||
|
||||
try:
|
||||
with connection.cursor() as cursor:
|
||||
cursor.execute(reviewed_query())
|
||||
columns = [description.name for description in cursor.description]
|
||||
rows = cursor.fetchall()
|
||||
connection.rollback()
|
||||
except psycopg2.Error as exc:
|
||||
connection.rollback()
|
||||
print(
|
||||
"Aggregate outcome audit failed: "
|
||||
f"{exc.diag.message_primary or type(exc).__name__}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return 1
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
widths = [len(column) for column in columns]
|
||||
rendered = [["" if value is None else str(value) for value in row] for row in rows]
|
||||
for row in rendered:
|
||||
widths = [max(width, len(value)) for width, value in zip(widths, row)]
|
||||
print(" ".join(value.ljust(width) for value, width in zip(columns, widths)))
|
||||
print(" ".join("-" * width for width in widths))
|
||||
for row in rendered:
|
||||
print(" ".join(value.ljust(width) for value, width in zip(row, widths)))
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,799 @@
|
||||
"""Train a leakage-safe nonlinear candidate from the private episode mart.
|
||||
|
||||
The model is fitted only on non-audit 2016-2022 episodes, selected on 2023,
|
||||
and probability-calibrated on 2024. The locked 2025 partition is evaluated
|
||||
only when ``--evaluate-locked`` is supplied. No row-level predictions are
|
||||
written.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import statistics
|
||||
import sys
|
||||
from collections import Counter
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Mapping, Optional, Sequence, Tuple
|
||||
|
||||
try:
|
||||
from scripts import train_baselines as baselines
|
||||
except ModuleNotFoundError: # Direct execution places scripts/ on sys.path.
|
||||
import train_baselines as baselines
|
||||
|
||||
|
||||
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
||||
MODEL_VERSION = "hist_gradient_boosting_v1"
|
||||
DEFAULT_OUTPUT = PROJECT_ROOT / "artifacts/private/tree" / MODEL_VERSION
|
||||
|
||||
NUMERIC_FEATURES = baselines.NUMERIC_FEATURES
|
||||
CATEGORICAL_FEATURES = baselines.CATEGORICAL_FEATURES
|
||||
|
||||
MISSING_CATEGORY_CODE = 0
|
||||
RARE_CATEGORY_CODE = 1
|
||||
UNKNOWN_CATEGORY_CODE = 2
|
||||
FIRST_KNOWN_CATEGORY_CODE = 3
|
||||
MAX_HISTOGRAM_BINS = 255
|
||||
|
||||
# These columns are validated by the baseline mart loader but are deliberately
|
||||
# unavailable to the nonlinear feature encoder.
|
||||
EXCLUDED_FROM_PREDICTORS = (
|
||||
"vehicle_token",
|
||||
"vehicle_bucket",
|
||||
"is_vin_audit",
|
||||
"episode_number",
|
||||
"episode_start",
|
||||
"first_outcome",
|
||||
"target_nonpass",
|
||||
"eligible_returning_target",
|
||||
"temporal_partition",
|
||||
"source_era",
|
||||
"target_outcome_label_source",
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TreeCandidate:
|
||||
learning_rate: float
|
||||
max_leaf_nodes: int
|
||||
l2_regularization: float
|
||||
|
||||
def artifact_state(self) -> Dict[str, object]:
|
||||
return {
|
||||
"learning_rate": self.learning_rate,
|
||||
"max_leaf_nodes": self.max_leaf_nodes,
|
||||
"l2_regularization": self.l2_regularization,
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class TreeFeatureEncoder:
|
||||
"""Train-only median imputation and bounded ordinal category encoding."""
|
||||
|
||||
min_category_count: int
|
||||
max_categories: int
|
||||
numeric_medians: Dict[str, float]
|
||||
seen_categories: Dict[str, set]
|
||||
known_category_codes: Dict[str, Dict[str, int]]
|
||||
feature_names: Tuple[str, ...]
|
||||
categorical_mask: Tuple[bool, ...]
|
||||
|
||||
@classmethod
|
||||
def fit(
|
||||
cls,
|
||||
rows: Sequence[baselines.EpisodeRow],
|
||||
min_category_count: int,
|
||||
max_categories: int,
|
||||
) -> "TreeFeatureEncoder":
|
||||
if not rows:
|
||||
raise baselines.DataValidationError(
|
||||
"Cannot fit tree preprocessing on zero rows"
|
||||
)
|
||||
if min_category_count < 1:
|
||||
raise baselines.DataValidationError(
|
||||
"min_category_count must be at least 1"
|
||||
)
|
||||
if not FIRST_KNOWN_CATEGORY_CODE + 1 <= max_categories <= MAX_HISTOGRAM_BINS:
|
||||
raise baselines.DataValidationError(
|
||||
"max_categories must be between {} and {}".format(
|
||||
FIRST_KNOWN_CATEGORY_CODE + 1, MAX_HISTOGRAM_BINS
|
||||
)
|
||||
)
|
||||
|
||||
numeric_medians: Dict[str, float] = {}
|
||||
for name in NUMERIC_FEATURES:
|
||||
observed = [
|
||||
row.numeric[name]
|
||||
for row in rows
|
||||
if row.numeric[name] is not None
|
||||
]
|
||||
if not observed:
|
||||
raise baselines.DataValidationError(
|
||||
"Training data has no observed values for numeric feature " + name
|
||||
)
|
||||
median = float(statistics.median(observed))
|
||||
if not math.isfinite(median):
|
||||
raise baselines.DataValidationError(
|
||||
"Training median is non-finite for numeric feature " + name
|
||||
)
|
||||
numeric_medians[name] = median
|
||||
|
||||
seen_categories: Dict[str, set] = {}
|
||||
known_category_codes: Dict[str, Dict[str, int]] = {}
|
||||
category_capacity = max_categories - FIRST_KNOWN_CATEGORY_CODE
|
||||
for name in CATEGORICAL_FEATURES:
|
||||
observed_values = [
|
||||
row.categorical[name]
|
||||
for row in rows
|
||||
if row.categorical[name] is not None
|
||||
]
|
||||
counts = Counter(observed_values)
|
||||
seen_categories[name] = set(counts)
|
||||
candidates = [
|
||||
(value, count)
|
||||
for value, count in counts.items()
|
||||
if count >= min_category_count
|
||||
]
|
||||
# Frequency first, lexical second makes capping deterministic.
|
||||
candidates.sort(key=lambda item: (-item[1], item[0]))
|
||||
kept = candidates[:category_capacity]
|
||||
known_category_codes[name] = {
|
||||
value: FIRST_KNOWN_CATEGORY_CODE + index
|
||||
for index, (value, _count) in enumerate(kept)
|
||||
}
|
||||
|
||||
feature_names = tuple(
|
||||
list(NUMERIC_FEATURES)
|
||||
+ ["missing__" + name for name in NUMERIC_FEATURES]
|
||||
+ list(CATEGORICAL_FEATURES)
|
||||
)
|
||||
categorical_mask = tuple(
|
||||
[False] * (2 * len(NUMERIC_FEATURES))
|
||||
+ [True] * len(CATEGORICAL_FEATURES)
|
||||
)
|
||||
return cls(
|
||||
min_category_count=min_category_count,
|
||||
max_categories=max_categories,
|
||||
numeric_medians=numeric_medians,
|
||||
seen_categories=seen_categories,
|
||||
known_category_codes=known_category_codes,
|
||||
feature_names=feature_names,
|
||||
categorical_mask=categorical_mask,
|
||||
)
|
||||
|
||||
def category_code(self, feature: str, value: Optional[str]) -> int:
|
||||
if value is None:
|
||||
return MISSING_CATEGORY_CODE
|
||||
known = self.known_category_codes[feature]
|
||||
if value in known:
|
||||
return known[value]
|
||||
if value in self.seen_categories[feature]:
|
||||
return RARE_CATEGORY_CODE
|
||||
return UNKNOWN_CATEGORY_CODE
|
||||
|
||||
def transform(self, rows: Sequence[baselines.EpisodeRow], np_module: object) -> object:
|
||||
matrix = np_module.empty((len(rows), len(self.feature_names)), dtype=np_module.float64)
|
||||
for row_index, row in enumerate(rows):
|
||||
column = 0
|
||||
for name in NUMERIC_FEATURES:
|
||||
value = row.numeric[name]
|
||||
matrix[row_index, column] = (
|
||||
self.numeric_medians[name] if value is None else float(value)
|
||||
)
|
||||
column += 1
|
||||
for name in NUMERIC_FEATURES:
|
||||
matrix[row_index, column] = 1.0 if row.numeric[name] is None else 0.0
|
||||
column += 1
|
||||
for name in CATEGORICAL_FEATURES:
|
||||
matrix[row_index, column] = float(
|
||||
self.category_code(name, row.categorical[name])
|
||||
)
|
||||
column += 1
|
||||
if matrix.size and not bool(np_module.isfinite(matrix).all()):
|
||||
raise baselines.DataValidationError(
|
||||
"Tree feature matrix contains non-finite values"
|
||||
)
|
||||
return matrix
|
||||
|
||||
def artifact_state(self) -> Dict[str, object]:
|
||||
return {
|
||||
"min_category_count": self.min_category_count,
|
||||
"max_categories": self.max_categories,
|
||||
"reserved_category_codes": {
|
||||
"missing": MISSING_CATEGORY_CODE,
|
||||
"rare_seen_in_training": RARE_CATEGORY_CODE,
|
||||
"unknown_after_training": UNKNOWN_CATEGORY_CODE,
|
||||
"first_known": FIRST_KNOWN_CATEGORY_CODE,
|
||||
},
|
||||
"numeric_medians": self.numeric_medians,
|
||||
"seen_categories": {
|
||||
name: sorted(values) for name, values in self.seen_categories.items()
|
||||
},
|
||||
"known_category_codes": self.known_category_codes,
|
||||
"feature_names": self.feature_names,
|
||||
"categorical_mask": self.categorical_mask,
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class TrainedTreeModel:
|
||||
encoder: TreeFeatureEncoder
|
||||
model: object
|
||||
platt_model: object
|
||||
selected_candidate: TreeCandidate
|
||||
model_n_iter: int
|
||||
platt_n_iter: int
|
||||
max_iter: int
|
||||
calibration_max_iter: int
|
||||
tuning_results: List[Dict[str, object]]
|
||||
|
||||
|
||||
def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--mart", required=True, type=Path)
|
||||
parser.add_argument("--output-dir", type=Path, default=DEFAULT_OUTPUT)
|
||||
parser.add_argument("--seed", type=int, default=20260715)
|
||||
parser.add_argument("--min-category-count", type=int, default=100)
|
||||
parser.add_argument("--max-categories", type=int, default=128)
|
||||
parser.add_argument("--learning-rates", default="0.05,0.1")
|
||||
parser.add_argument("--max-leaf-nodes", default="15,31")
|
||||
parser.add_argument("--l2-grid", default="1.0")
|
||||
parser.add_argument("--min-samples-leaf", type=int, default=20)
|
||||
parser.add_argument("--max-iter", type=int, default=300)
|
||||
parser.add_argument("--n-iter-no-change", type=int, default=20)
|
||||
parser.add_argument("--calibration-max-iter", type=int, default=5000)
|
||||
parser.add_argument("--calibration-bins", type=int, default=10)
|
||||
parser.add_argument("--evaluate-locked", action="store_true")
|
||||
parser.add_argument("--overwrite", action="store_true")
|
||||
return parser.parse_args(argv)
|
||||
|
||||
|
||||
def _parse_float_grid(value: str, label: str, allow_zero: bool) -> List[float]:
|
||||
try:
|
||||
values = [float(part.strip()) for part in value.split(",") if part.strip()]
|
||||
except ValueError as exc:
|
||||
raise baselines.DataValidationError(label + " must contain finite numbers") from exc
|
||||
minimum_ok = (lambda item: item >= 0.0) if allow_zero else (lambda item: item > 0.0)
|
||||
if not values or any(not math.isfinite(item) or not minimum_ok(item) for item in values):
|
||||
raise baselines.DataValidationError(label + " contains an invalid value")
|
||||
return sorted(set(values))
|
||||
|
||||
|
||||
def _parse_int_grid(value: str, label: str, minimum: int) -> List[int]:
|
||||
try:
|
||||
values = [int(part.strip()) for part in value.split(",") if part.strip()]
|
||||
except ValueError as exc:
|
||||
raise baselines.DataValidationError(label + " must contain integers") from exc
|
||||
if not values or any(item < minimum for item in values):
|
||||
raise baselines.DataValidationError(label + " contains an invalid value")
|
||||
return sorted(set(values))
|
||||
|
||||
|
||||
def candidate_grid(
|
||||
learning_rates: Sequence[float],
|
||||
max_leaf_nodes: Sequence[int],
|
||||
l2_values: Sequence[float],
|
||||
) -> List[TreeCandidate]:
|
||||
return [
|
||||
TreeCandidate(rate, leaves, l2)
|
||||
for rate in sorted(set(learning_rates))
|
||||
for leaves in sorted(set(max_leaf_nodes))
|
||||
for l2 in sorted(set(l2_values))
|
||||
]
|
||||
|
||||
|
||||
def require_tree_dependencies() -> Tuple[object, object, object, object, object, object]:
|
||||
# joblib's macOS physical-core probe emits a UserWarning on this managed
|
||||
# host. Preserve an explicit user setting; otherwise give it the already
|
||||
# available logical count so warning-as-error validation remains usable.
|
||||
os.environ.setdefault("LOKY_MAX_CPU_COUNT", "1")
|
||||
try:
|
||||
import joblib
|
||||
import numpy as np
|
||||
import sklearn
|
||||
from sklearn.ensemble import HistGradientBoostingClassifier
|
||||
from sklearn.exceptions import ConvergenceWarning
|
||||
from sklearn.linear_model import LogisticRegression
|
||||
except ImportError as exc:
|
||||
raise baselines.DependencyError(
|
||||
"Tree training requires numpy, scikit-learn, and joblib"
|
||||
) from exc
|
||||
return (
|
||||
np,
|
||||
sklearn,
|
||||
joblib,
|
||||
HistGradientBoostingClassifier,
|
||||
LogisticRegression,
|
||||
ConvergenceWarning,
|
||||
)
|
||||
|
||||
|
||||
def _non_audit(rows: Sequence[baselines.EpisodeRow]) -> List[baselines.EpisodeRow]:
|
||||
return [row for row in rows if not row.audit_vehicle]
|
||||
|
||||
|
||||
def _targets(rows: Sequence[baselines.EpisodeRow]) -> List[int]:
|
||||
return [row.target for row in rows]
|
||||
|
||||
|
||||
def _scores_finite(model: object, np_module: object) -> bool:
|
||||
for attribute in ("train_score_", "validation_score_"):
|
||||
values = getattr(model, attribute, None)
|
||||
if values is None:
|
||||
continue
|
||||
array = np_module.asarray(values)
|
||||
if array.size and not bool(np_module.isfinite(array).all()):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _hist_converged(
|
||||
model: object,
|
||||
convergence_messages: Sequence[str],
|
||||
max_iter: int,
|
||||
np_module: object,
|
||||
) -> Tuple[bool, Optional[int]]:
|
||||
value = getattr(model, "n_iter_", None)
|
||||
if value is None:
|
||||
return False, None
|
||||
try:
|
||||
n_iter = int(value)
|
||||
except (TypeError, ValueError, OverflowError):
|
||||
return False, None
|
||||
converged = (
|
||||
not convergence_messages
|
||||
and 0 < n_iter < max_iter
|
||||
and _scores_finite(model, np_module)
|
||||
)
|
||||
return converged, n_iter
|
||||
|
||||
|
||||
def _decision_scores(model: object, features: object, np_module: object) -> object:
|
||||
scores = np_module.asarray(model.decision_function(features), dtype=np_module.float64)
|
||||
scores = scores.reshape(-1, 1)
|
||||
if scores.size == 0 or not bool(np_module.isfinite(scores).all()):
|
||||
raise baselines.DataValidationError(
|
||||
"Tree model produced non-finite decision scores"
|
||||
)
|
||||
return scores
|
||||
|
||||
|
||||
def train_tree_model(
|
||||
mart: baselines.MartData,
|
||||
candidates: Sequence[TreeCandidate],
|
||||
min_category_count: int,
|
||||
max_categories: int,
|
||||
min_samples_leaf: int,
|
||||
max_iter: int,
|
||||
n_iter_no_change: int,
|
||||
calibration_max_iter: int,
|
||||
seed: int,
|
||||
) -> TrainedTreeModel:
|
||||
(
|
||||
np,
|
||||
_sklearn,
|
||||
_joblib,
|
||||
hist_class,
|
||||
logistic_class,
|
||||
convergence_warning_class,
|
||||
) = require_tree_dependencies()
|
||||
if not candidates:
|
||||
raise baselines.DataValidationError("The tree tuning grid is empty")
|
||||
if min_samples_leaf < 1 or max_iter < 2 or n_iter_no_change < 1:
|
||||
raise baselines.DataValidationError("Invalid tree iteration/leaf configuration")
|
||||
if calibration_max_iter < 2:
|
||||
raise baselines.DataValidationError("calibration_max_iter must be at least 2")
|
||||
|
||||
train_rows = _non_audit(mart.rows_by_partition["train"])
|
||||
tune_rows = _non_audit(mart.rows_by_partition["tune"])
|
||||
calibrate_rows = _non_audit(mart.rows_by_partition["calibrate"])
|
||||
baselines._require_two_classes(train_rows, "Tree training")
|
||||
baselines._require_two_classes(tune_rows, "Tree tuning")
|
||||
baselines._require_two_classes(calibrate_rows, "Tree calibration")
|
||||
|
||||
encoder = TreeFeatureEncoder.fit(
|
||||
train_rows,
|
||||
min_category_count=min_category_count,
|
||||
max_categories=max_categories,
|
||||
)
|
||||
x_train = encoder.transform(train_rows, np)
|
||||
x_tune = encoder.transform(tune_rows, np)
|
||||
x_calibrate = encoder.transform(calibrate_rows, np)
|
||||
train_targets = np.asarray(_targets(train_rows), dtype=np.int8)
|
||||
tune_targets = _targets(tune_rows)
|
||||
calibrate_targets = np.asarray(_targets(calibrate_rows), dtype=np.int8)
|
||||
|
||||
selectable: List[Tuple[float, float, float, int, float, object, int, TreeCandidate]] = []
|
||||
tuning_results: List[Dict[str, object]] = []
|
||||
for candidate in candidates:
|
||||
model = hist_class(
|
||||
loss="log_loss",
|
||||
learning_rate=candidate.learning_rate,
|
||||
max_iter=max_iter,
|
||||
max_leaf_nodes=candidate.max_leaf_nodes,
|
||||
min_samples_leaf=min_samples_leaf,
|
||||
l2_regularization=candidate.l2_regularization,
|
||||
max_bins=MAX_HISTOGRAM_BINS,
|
||||
categorical_features=list(encoder.categorical_mask),
|
||||
early_stopping=True,
|
||||
scoring="loss",
|
||||
validation_fraction=0.10,
|
||||
n_iter_no_change=n_iter_no_change,
|
||||
tol=1e-7,
|
||||
random_state=seed,
|
||||
class_weight=None,
|
||||
)
|
||||
result: Dict[str, object] = {
|
||||
**candidate.artifact_state(),
|
||||
"converged": False,
|
||||
"n_iter": None,
|
||||
"convergence_warning": None,
|
||||
"finite_scores": False,
|
||||
"finite_train_probabilities": False,
|
||||
"finite_tune_probabilities": False,
|
||||
"eligible_for_selection": False,
|
||||
"brier": None,
|
||||
"average_precision": None,
|
||||
}
|
||||
convergence_messages = baselines.fit_with_convergence_capture(
|
||||
model, x_train, train_targets, convergence_warning_class
|
||||
)
|
||||
converged, n_iter = _hist_converged(
|
||||
model, convergence_messages, max_iter, np
|
||||
)
|
||||
result["converged"] = converged
|
||||
result["n_iter"] = n_iter
|
||||
result["convergence_warning"] = (
|
||||
" | ".join(convergence_messages) if convergence_messages else None
|
||||
)
|
||||
result["finite_scores"] = _scores_finite(model, np)
|
||||
|
||||
train_probabilities = model.predict_proba(x_train)[:, 1]
|
||||
tune_probabilities = model.predict_proba(x_tune)[:, 1]
|
||||
finite_train = baselines.probability_array_finite(train_probabilities, np)
|
||||
finite_tune = baselines.probability_array_finite(tune_probabilities, np)
|
||||
result["finite_train_probabilities"] = finite_train
|
||||
result["finite_tune_probabilities"] = finite_tune
|
||||
eligible = converged and finite_train and finite_tune
|
||||
result["eligible_for_selection"] = eligible
|
||||
if eligible and n_iter is not None:
|
||||
metrics = baselines.binary_metrics(tune_targets, tune_probabilities.tolist())
|
||||
brier = float(metrics["brier"])
|
||||
average_precision = float(metrics["average_precision"])
|
||||
result["brier"] = brier
|
||||
result["average_precision"] = average_precision
|
||||
selectable.append(
|
||||
(
|
||||
brier,
|
||||
-average_precision,
|
||||
candidate.learning_rate,
|
||||
candidate.max_leaf_nodes,
|
||||
candidate.l2_regularization,
|
||||
model,
|
||||
n_iter,
|
||||
candidate,
|
||||
)
|
||||
)
|
||||
tuning_results.append(result)
|
||||
|
||||
if not selectable:
|
||||
raise baselines.DataValidationError(
|
||||
"No tree candidate converged with finite train/tune probabilities; "
|
||||
"increase --max-iter or inspect the feature contract"
|
||||
)
|
||||
(
|
||||
_brier,
|
||||
_negative_ap,
|
||||
_rate,
|
||||
_leaves,
|
||||
_l2,
|
||||
selected_model,
|
||||
selected_n_iter,
|
||||
selected_candidate,
|
||||
) = min(selectable, key=lambda item: item[:5])
|
||||
|
||||
calibration_scores = _decision_scores(selected_model, x_calibrate, np)
|
||||
platt_model = logistic_class(
|
||||
C=baselines.PLATT_CALIBRATION_CONFIG["c"],
|
||||
penalty=baselines.PLATT_CALIBRATION_CONFIG["penalty"],
|
||||
solver=baselines.PLATT_CALIBRATION_CONFIG["solver"],
|
||||
class_weight=None,
|
||||
max_iter=calibration_max_iter,
|
||||
random_state=seed,
|
||||
)
|
||||
platt_warnings = baselines.fit_with_convergence_capture(
|
||||
platt_model,
|
||||
calibration_scores,
|
||||
calibrate_targets,
|
||||
convergence_warning_class,
|
||||
)
|
||||
platt_n_iter = baselines.model_n_iter(platt_model, np)
|
||||
if (
|
||||
platt_warnings
|
||||
or platt_n_iter is None
|
||||
or platt_n_iter >= calibration_max_iter
|
||||
or not baselines.model_parameters_finite(platt_model, np)
|
||||
):
|
||||
detail = " | ".join(platt_warnings) if platt_warnings else "finite/convergence check"
|
||||
raise baselines.DataValidationError(
|
||||
"Tree Platt calibration did not converge: " + detail
|
||||
)
|
||||
calibrated = platt_model.predict_proba(calibration_scores)[:, 1]
|
||||
if not baselines.probability_array_finite(calibrated, np):
|
||||
raise baselines.DataValidationError(
|
||||
"Tree Platt calibration produced non-finite probabilities"
|
||||
)
|
||||
|
||||
return TrainedTreeModel(
|
||||
encoder=encoder,
|
||||
model=selected_model,
|
||||
platt_model=platt_model,
|
||||
selected_candidate=selected_candidate,
|
||||
model_n_iter=selected_n_iter,
|
||||
platt_n_iter=int(platt_n_iter),
|
||||
max_iter=max_iter,
|
||||
calibration_max_iter=calibration_max_iter,
|
||||
tuning_results=tuning_results,
|
||||
)
|
||||
|
||||
|
||||
def tree_probabilities(
|
||||
trained: TrainedTreeModel,
|
||||
rows: Sequence[baselines.EpisodeRow],
|
||||
calibrated: bool,
|
||||
) -> List[float]:
|
||||
np, _sklearn, _joblib, _hist, _logistic, _warning = require_tree_dependencies()
|
||||
matrix = trained.encoder.transform(rows, np)
|
||||
if calibrated:
|
||||
scores = _decision_scores(trained.model, matrix, np)
|
||||
probabilities = trained.platt_model.predict_proba(scores)[:, 1]
|
||||
else:
|
||||
probabilities = trained.model.predict_proba(matrix)[:, 1]
|
||||
if not baselines.probability_array_finite(probabilities, np):
|
||||
raise baselines.DataValidationError(
|
||||
"Tree model produced non-finite probabilities"
|
||||
)
|
||||
return probabilities.tolist()
|
||||
|
||||
|
||||
def evaluate_tree_model(
|
||||
mart: baselines.MartData,
|
||||
trained: TrainedTreeModel,
|
||||
evaluate_locked: bool,
|
||||
bin_count: int,
|
||||
) -> Tuple[List[Dict[str, object]], List[Dict[str, object]]]:
|
||||
metric_rows: List[Dict[str, object]] = []
|
||||
calibration_rows: List[Dict[str, object]] = []
|
||||
partitions = ["train", "tune", "calibrate"]
|
||||
if evaluate_locked:
|
||||
partitions.append("locked_test")
|
||||
|
||||
for partition in partitions:
|
||||
source_rows = mart.rows_by_partition[partition]
|
||||
cohorts = baselines.evaluation_cohorts(
|
||||
source_rows, include_audit_breakout=(partition == "locked_test")
|
||||
)
|
||||
for cohort_name, rows in cohorts:
|
||||
if not rows:
|
||||
continue
|
||||
targets = _targets(rows)
|
||||
predictions = [
|
||||
(
|
||||
"hist_gradient_boosting_raw",
|
||||
tree_probabilities(trained, rows, calibrated=False),
|
||||
)
|
||||
]
|
||||
if partition in {"calibrate", "locked_test"}:
|
||||
predictions.append(
|
||||
(
|
||||
"hist_gradient_boosting_platt",
|
||||
tree_probabilities(trained, rows, calibrated=True),
|
||||
)
|
||||
)
|
||||
for model_name, probabilities in predictions:
|
||||
metrics = baselines.binary_metrics(targets, probabilities)
|
||||
metric_rows.append(
|
||||
{
|
||||
"model": model_name,
|
||||
"partition": partition,
|
||||
"cohort": cohort_name,
|
||||
"episodes": len(rows),
|
||||
"vehicles": len({row.vehicle_token for row in rows}),
|
||||
"nonpass": sum(targets),
|
||||
**metrics,
|
||||
}
|
||||
)
|
||||
for values in baselines.calibration_bins(
|
||||
targets, probabilities, bin_count
|
||||
):
|
||||
calibration_rows.append(
|
||||
{
|
||||
"model": model_name,
|
||||
"partition": partition,
|
||||
"cohort": cohort_name,
|
||||
**values,
|
||||
}
|
||||
)
|
||||
return metric_rows, calibration_rows
|
||||
|
||||
|
||||
def _input_sha256(path: Path) -> str:
|
||||
digest = hashlib.sha256()
|
||||
with path.open("rb") as handle:
|
||||
for block in iter(lambda: handle.read(1024 * 1024), b""):
|
||||
digest.update(block)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def write_artifacts(
|
||||
output_dir: Path,
|
||||
mart_path: Path,
|
||||
mart: baselines.MartData,
|
||||
trained: TrainedTreeModel,
|
||||
metric_rows: Sequence[Mapping[str, object]],
|
||||
calibration_rows: Sequence[Mapping[str, object]],
|
||||
evaluate_locked: bool,
|
||||
seed: int,
|
||||
overwrite: bool,
|
||||
sklearn_version: str,
|
||||
) -> None:
|
||||
_np, _sklearn, joblib, _hist, _logistic, _warning = require_tree_dependencies()
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
model_path = output_dir / "model.joblib"
|
||||
metrics_path = output_dir / "metrics.json"
|
||||
calibration_path = output_dir / "calibration_bins.csv"
|
||||
manifest_path = output_dir / "manifest.json"
|
||||
expected = (model_path, metrics_path, calibration_path, manifest_path)
|
||||
existing = [path for path in expected if path.exists()]
|
||||
if existing and not overwrite:
|
||||
raise baselines.DataValidationError(
|
||||
"Refusing to overwrite existing tree artifacts: "
|
||||
+ ", ".join(path.name for path in existing)
|
||||
)
|
||||
|
||||
partial_model = model_path.with_name(model_path.name + ".partial")
|
||||
joblib.dump(
|
||||
{
|
||||
"model_version": MODEL_VERSION,
|
||||
"target_contract": baselines.TARGET_CONTRACT,
|
||||
"numeric_features": NUMERIC_FEATURES,
|
||||
"categorical_features": CATEGORICAL_FEATURES,
|
||||
"excluded_from_predictors": EXCLUDED_FROM_PREDICTORS,
|
||||
"encoder": trained.encoder.artifact_state(),
|
||||
"selected_candidate": trained.selected_candidate.artifact_state(),
|
||||
"hist_gradient_boosting_model": trained.model,
|
||||
"platt_model": trained.platt_model,
|
||||
"platt_calibration_config": baselines.PLATT_CALIBRATION_CONFIG,
|
||||
},
|
||||
partial_model,
|
||||
)
|
||||
os.replace(partial_model, model_path)
|
||||
baselines._atomic_text(
|
||||
metrics_path,
|
||||
json.dumps(list(metric_rows), indent=2, sort_keys=True, allow_nan=False)
|
||||
+ "\n",
|
||||
)
|
||||
baselines._atomic_csv(calibration_path, calibration_rows)
|
||||
|
||||
partition_counts, source_counts, label_source_counts = (
|
||||
baselines.manifest_audit_counts(mart, evaluate_locked=evaluate_locked)
|
||||
)
|
||||
manifest = {
|
||||
"model_version": MODEL_VERSION,
|
||||
"generated_at_utc": datetime.now(timezone.utc).isoformat(),
|
||||
"classification": "private_model_artifact_no_row_predictions",
|
||||
"input_file": mart_path.name,
|
||||
"input_sha256": _input_sha256(mart_path),
|
||||
"metrics_sha256": _input_sha256(metrics_path),
|
||||
"input_rows": mart.input_rows,
|
||||
"eligible_rows": mart.eligible_rows,
|
||||
"partition_counts": partition_counts,
|
||||
"source_era_audit_counts": source_counts,
|
||||
"target_label_source_counts": label_source_counts,
|
||||
"target_contract": baselines.TARGET_CONTRACT,
|
||||
"never_fit_audit_rule": "mart is_vin_audit; validated as vehicle_bucket < 10",
|
||||
"locked_test_evaluated": evaluate_locked,
|
||||
"split_bounds": {
|
||||
name: {
|
||||
"start_inclusive": start.isoformat(),
|
||||
"end_exclusive": end.isoformat(),
|
||||
}
|
||||
for name, (start, end) in baselines.SPLIT_BOUNDS.items()
|
||||
if name != "shadow"
|
||||
},
|
||||
"numeric_features": NUMERIC_FEATURES,
|
||||
"categorical_features": CATEGORICAL_FEATURES,
|
||||
"excluded_from_predictors": EXCLUDED_FROM_PREDICTORS,
|
||||
"preprocessing_fit_partition": "train_2016_2022_non_audit_only",
|
||||
"unknown_category_code": UNKNOWN_CATEGORY_CODE,
|
||||
"selected_candidate": trained.selected_candidate.artifact_state(),
|
||||
"convergence": {
|
||||
"tree_max_iter": trained.max_iter,
|
||||
"tree_selected_n_iter": trained.model_n_iter,
|
||||
"platt_max_iter": trained.calibration_max_iter,
|
||||
"platt_n_iter": trained.platt_n_iter,
|
||||
"requires_early_stop_before_max_iter": True,
|
||||
"requires_finite_scores_and_probabilities": True,
|
||||
},
|
||||
"platt_calibration_config": baselines.PLATT_CALIBRATION_CONFIG,
|
||||
"tuning_results": trained.tuning_results,
|
||||
"seed": seed,
|
||||
"python_version": sys.version,
|
||||
"scikit_learn_version": sklearn_version,
|
||||
"artifacts": [path.name for path in expected],
|
||||
}
|
||||
baselines._atomic_text(
|
||||
manifest_path,
|
||||
json.dumps(manifest, indent=2, sort_keys=True, allow_nan=False) + "\n",
|
||||
)
|
||||
|
||||
|
||||
def main(argv: Optional[Sequence[str]] = None) -> int:
|
||||
args = parse_args(argv)
|
||||
os.umask(0o077)
|
||||
try:
|
||||
mart_path = baselines.require_private_path(args.mart, "--mart")
|
||||
output_dir = baselines.require_private_path(args.output_dir, "--output-dir")
|
||||
learning_rates = _parse_float_grid(
|
||||
args.learning_rates, "--learning-rates", allow_zero=False
|
||||
)
|
||||
leaf_nodes = _parse_int_grid(
|
||||
args.max_leaf_nodes, "--max-leaf-nodes", minimum=2
|
||||
)
|
||||
l2_values = _parse_float_grid(args.l2_grid, "--l2-grid", allow_zero=True)
|
||||
if args.calibration_bins < 2:
|
||||
raise baselines.DataValidationError(
|
||||
"--calibration-bins must be at least 2"
|
||||
)
|
||||
mart = baselines.load_mart(mart_path)
|
||||
trained = train_tree_model(
|
||||
mart=mart,
|
||||
candidates=candidate_grid(learning_rates, leaf_nodes, l2_values),
|
||||
min_category_count=args.min_category_count,
|
||||
max_categories=args.max_categories,
|
||||
min_samples_leaf=args.min_samples_leaf,
|
||||
max_iter=args.max_iter,
|
||||
n_iter_no_change=args.n_iter_no_change,
|
||||
calibration_max_iter=args.calibration_max_iter,
|
||||
seed=args.seed,
|
||||
)
|
||||
metric_rows, calibration_rows = evaluate_tree_model(
|
||||
mart,
|
||||
trained,
|
||||
evaluate_locked=args.evaluate_locked,
|
||||
bin_count=args.calibration_bins,
|
||||
)
|
||||
_np, sklearn, _joblib, _hist, _logistic, _warning = (
|
||||
require_tree_dependencies()
|
||||
)
|
||||
write_artifacts(
|
||||
output_dir=output_dir,
|
||||
mart_path=mart_path,
|
||||
mart=mart,
|
||||
trained=trained,
|
||||
metric_rows=metric_rows,
|
||||
calibration_rows=calibration_rows,
|
||||
evaluate_locked=args.evaluate_locked,
|
||||
seed=args.seed,
|
||||
overwrite=args.overwrite,
|
||||
sklearn_version=sklearn.__version__,
|
||||
)
|
||||
except baselines.BaselineError as exc:
|
||||
print("Tree training failed: {}".format(exc), file=sys.stderr)
|
||||
return 2
|
||||
|
||||
print("Wrote private nonlinear model artifacts to {}".format(output_dir))
|
||||
if args.evaluate_locked:
|
||||
print("The explicitly unlocked 2025 test metrics were evaluated.")
|
||||
else:
|
||||
print("The 2025 locked test was not evaluated.")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
Reference in New Issue
Block a user