#!/usr/bin/env python
"""
Risk Table Mapping Processing Script
This script creates and applies risk tables for categorical features based on
target variable correlation, and handles missing value imputation for numeric features.
It supports both training mode (fit and transform) and inference mode (transform only).
"""
import argparse
import os
import sys
import pandas as pd
import numpy as np
import json
import pickle as pkl
import traceback
import shutil
import gc
from pathlib import Path
from collections import Counter
from multiprocessing import Pool, cpu_count
from sklearn.impute import SimpleImputer
import logging
from typing import Dict, List, Tuple, Any, Optional, Callable
# Default paths (will be overridden by parameters in main function)
DEFAULT_INPUT_DIR = "/opt/ml/processing/input/data"
DEFAULT_CONFIG_DIR = "/opt/ml/code/hyperparams" # Source directory path
DEFAULT_OUTPUT_DIR = "/opt/ml/processing/output"
DEFAULT_MODEL_ARTIFACTS_DIR = "/opt/ml/processing/input/model_artifacts"
# Constants for file paths to ensure consistency between training and inference
# Match XGBoost training output format
RISK_TABLE_FILENAME = "risk_table_map.pkl" # Matches XGBoost training output
HYPERPARAMS_FILENAME = (
"hyperparameters.json" # Expected by the script and generated by the builder
)
# Set up logging
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s - %(levelname)s - [%(filename)s:%(lineno)d] - %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger(__name__)
# ============================================================================
# STREAMING MODE UTILITIES (Reused from missing_value_imputation)
# ============================================================================
[docs]
def collect_risk_table_statistics_pass1(
all_shards: List[Path],
signature_columns: Optional[List[str]],
cat_field_list: List[str],
label_name: str,
smooth_factor: float,
count_threshold: int,
max_unique_threshold: int,
log_func: Callable,
) -> Dict[str, Any]:
"""
Pass 1: Collect risk table statistics from training shards.
Memory-efficient incremental crosstab aggregation:
- For each categorical field:
- Accumulate counts per (category, label) combination
- Compute risk scores after all shards processed
This function processes shards sequentially to build risk tables incrementally,
avoiding loading the entire dataset into memory.
Args:
all_shards: List of training shard paths
signature_columns: Optional column names for CSV/TSV
cat_field_list: List of categorical field names
label_name: Target variable name
smooth_factor: Smoothing factor for risk calculation
count_threshold: Minimum count threshold
max_unique_threshold: Max unique values for validation
log_func: Logging function
Returns:
Dictionary of risk tables (same format as OfflineBinning.risk_tables)
Format: {field: {varName, type, mode, default_bin, bins: {category: risk}}}
"""
log_func("[PASS1] Collecting risk table statistics from training shards...")
# ========================================================================
# STEP 1: Identify valid categorical fields from first shard
# ========================================================================
first_shard = all_shards[0]
df_first = _read_file_to_df(first_shard, signature_columns)
df_first.columns = [col.replace("__DOT__", ".") for col in df_first.columns]
# Validate categorical fields
valid_cat_fields = []
for field in cat_field_list:
if field in df_first.columns:
unique_count = df_first[field].nunique()
if unique_count < max_unique_threshold:
valid_cat_fields.append(field)
log_func(f"[PASS1] {field}: {unique_count} unique values")
else:
log_func(
f"[PASS1] {field}: SKIP ({unique_count} > {max_unique_threshold})"
)
if not valid_cat_fields:
log_func("[PASS1] No valid categorical fields found")
return {}
# ========================================================================
# STEP 2: Initialize crosstab accumulators
# ========================================================================
# For each field, accumulate: {(category, label): count}
crosstab_accumulators = {}
for field in valid_cat_fields:
crosstab_accumulators[field] = {
"counts": {}, # {(category, label): count}
"data_mode": None,
}
# Compute default risk from first shard (rough approximation)
df_fit = df_first.loc[
(df_first[label_name] != -1) & (~df_first[label_name].isnull())
].copy()
default_risk = float(df_fit[label_name].mean()) if len(df_fit) > 0 else 0.0
del df_first, df_fit
gc.collect()
# ========================================================================
# STEP 3: Process each shard and accumulate crosstabs
# ========================================================================
log_func(f"[PASS1] Processing {len(all_shards)} training shards...")
for i, shard_path in enumerate(all_shards):
try:
df = _read_file_to_df(shard_path, signature_columns)
df.columns = [col.replace("__DOT__", ".") for col in df.columns]
# Filter valid rows for fitting (exclude -1 and NaN labels)
df_fit = df.loc[(df[label_name] != -1) & (~df[label_name].isnull())].copy()
if len(df_fit) == 0:
continue
# Accumulate crosstabs for each field
for field in valid_cat_fields:
if field not in df_fit.columns:
continue
accumulator = crosstab_accumulators[field]
# Detect data mode from first shard with valid data
if accumulator["data_mode"] is None:
if pd.api.types.is_numeric_dtype(df_fit[field]):
accumulator["data_mode"] = "numeric"
else:
accumulator["data_mode"] = "categorical"
# Accumulate counts per (category, label) pair
for category, label in zip(df_fit[field], df_fit[label_name]):
if pd.isna(category):
continue
key = (category, label)
accumulator["counts"][key] = accumulator["counts"].get(key, 0) + 1
del df, df_fit
gc.collect()
# Progress logging every 100 shards
if (i + 1) % 100 == 0:
log_func(f"[PASS1] Processed {i + 1}/{len(all_shards)} shards")
except Exception as e:
log_func(f"[PASS1 WARNING] Failed to read {shard_path.name}: {e}")
continue
# ========================================================================
# STEP 4: Compute risk tables from accumulated crosstabs
# ========================================================================
log_func("[PASS1] Computing final risk tables...")
# Estimate smooth_samples based on dataset size
smooth_samples = int(len(all_shards) * 100 * smooth_factor) # Rough estimate
risk_tables = {}
for field, accumulator in crosstab_accumulators.items():
try:
counts = accumulator["counts"]
if not counts:
log_func(f"[PASS1] {field}: No valid data, using default")
risk_tables[field] = {
"varName": field,
"type": "categorical",
"mode": accumulator["data_mode"] or "categorical",
"default_bin": default_risk,
"bins": {},
}
continue
# Build risk table from accumulated counts
# Group by category, compute risk = count(label=1) / count(total)
category_stats = {}
for (category, label), count in counts.items():
if category not in category_stats:
category_stats[category] = {"total": 0, "positive": 0}
category_stats[category]["total"] += count
if label == 1:
category_stats[category]["positive"] += count
# Compute smoothed risk scores
bins = {}
for category, stats in category_stats.items():
total = stats["total"]
positive = stats["positive"]
# Base risk
risk = positive / total if total > 0 else 0.0
# Apply smoothing and threshold
if total >= count_threshold:
smooth_risk = (total * risk + smooth_samples * default_risk) / (
total + smooth_samples
)
else:
smooth_risk = default_risk
bins[category] = float(smooth_risk)
risk_tables[field] = {
"varName": field,
"type": "categorical",
"mode": accumulator["data_mode"],
"default_bin": default_risk,
"bins": bins,
}
log_func(f"[PASS1] {field}: {len(bins)} categories mapped")
except Exception as e:
log_func(f"[PASS1 WARNING] Failed to compute risk table for {field}: {e}")
risk_tables[field] = {
"varName": field,
"type": "categorical",
"mode": accumulator["data_mode"] or "categorical",
"default_bin": default_risk,
"bins": {},
}
log_func(f"[PASS1] Complete! Created risk tables for {len(risk_tables)} fields")
return risk_tables
[docs]
def process_shard_end_to_end_risk_mapping(args: tuple) -> Dict[str, int]:
"""
Process single shard: read → apply risk tables → write.
Stateless per-shard processing using global risk_tables from Pass 1.
Preserves 1:1 shard mapping (input shard number → output shard number).
This function is designed to be called in parallel via multiprocessing.Pool.
Each worker processes one shard completely independently.
Args:
args: Tuple of (shard_path, shard_num, global_context,
output_base, signature_columns, output_format)
global_context must contain:
- "risk_tables": Dictionary of risk tables from Pass 1
- "split_name": Which split this shard belongs to ("train", "val", "test", etc.)
Returns:
Statistics dict with row count for this split
Format: {"train": 1000} or {"val": 200} or {"validation": 500}
Example:
Input: train/part-00042.csv (1000 rows)
Output: train/part-00042.csv (1000 rows, risk-mapped)
Return: {"train": 1000}
"""
(
shard_path,
shard_num,
global_context,
output_base,
signature_columns,
output_format,
) = args
try:
# ====================================================================
# STEP 1: Read Single Shard
# ====================================================================
df = _read_file_to_df(shard_path, signature_columns)
df.columns = [col.replace("__DOT__", ".") for col in df.columns]
# ====================================================================
# STEP 2: Apply Risk Table Mapping (Using Global Context)
# ====================================================================
risk_tables = global_context["risk_tables"]
# Simple map operation for each categorical field
# This is fast because risk_tables are already computed
for field, risk_table_info in risk_tables.items():
if field in df.columns:
bins = risk_table_info["bins"]
default_bin = risk_table_info["default_bin"]
# Map categorical values to risk scores
df[field] = df[field].map(bins).fillna(default_bin)
# ====================================================================
# STEP 3: Write to Correct Split Folder (Preserving Shard Number)
# ====================================================================
split_name = global_context["split_name"]
stats = {}
if len(df) > 0:
# Preserve exact shard numbering: part-00042 → part-00042
output_path = (
output_base / split_name / f"part-{shard_num:05d}.{output_format}"
)
write_shard_file(df, output_path, output_format)
stats[split_name] = len(df)
else:
# Empty shard - skip writing (sparse numbering is OK)
stats[split_name] = 0
return stats
except Exception as e:
# Log error but don't crash the entire pool
# This allows other shards to continue processing
print(f"[ERROR] Failed to process shard {shard_num} ({shard_path.name}): {e}")
import traceback
print(traceback.format_exc())
# Return zero stats for this shard
split_name = global_context.get("split_name", "unknown")
return {split_name: 0}
[docs]
def process_streaming_mode_risk_mapping(
input_dir: str,
output_dir: str,
signature_columns: Optional[List[str]],
job_type: str,
hyperparams: Dict[str, Any],
environ_vars: Dict[str, str],
max_workers: Optional[int],
model_artifacts_input_dir: Optional[str] = None,
model_artifacts_output_dir: Optional[str] = None,
log_func: Optional[Callable] = None,
) -> Dict[str, int]:
"""
Streaming mode for risk table mapping with train/val/test subdirectories.
Two-pass architecture:
- Pass 1 (Sequential): Collect risk table statistics from training shards only
- Pass 2 (Parallel): Apply risk tables per split in parallel
Auto-detects output format from input shards (mirrors batch mode behavior).
Input structure (training mode):
input_dir/
train/part-00000.csv, part-00001.csv, ...
val/part-00000.csv, part-00001.csv, ...
test/part-00000.csv, part-00001.csv, ...
Output structure (training mode):
output_dir/
train/part-00000.csv, part-00001.csv, ... (risk-mapped, same format)
val/part-00000.csv, part-00001.csv, ... (risk-mapped, same format)
test/part-00000.csv, part-00001.csv, ... (risk-mapped, same format)
Args:
input_dir: Base input directory
output_dir: Base output directory
signature_columns: Optional column names for CSV/TSV
job_type: 'training', 'validation', 'testing', 'calibration'
hyperparams: Hyperparameters dict (contains cat_field_list, etc.)
environ_vars: Environment variables dict
max_workers: Number of parallel workers
model_artifacts_input_dir: Input model artifacts directory
model_artifacts_output_dir: Output model artifacts directory
log_func: Logging function (defaults to print)
Returns:
Dictionary with total row counts per split
Format: {"train": 100000, "val": 20000, "test": 30000}
"""
log = log_func or print
output_path = Path(output_dir)
# Extract parameters from hyperparameters
cat_field_list = hyperparams.get("cat_field_list", [])
label_name = hyperparams.get("label_name", "target")
smooth_factor = float(
environ_vars.get("SMOOTH_FACTOR", hyperparams.get("smooth_factor", 0.01))
)
count_threshold = int(
environ_vars.get("COUNT_THRESHOLD", hyperparams.get("count_threshold", 5))
)
max_unique_threshold = int(
environ_vars.get(
"MAX_UNIQUE_THRESHOLD", hyperparams.get("max_unique_threshold", 100)
)
)
# Determine optimal workers (0 or None = auto-detect)
if max_workers is None or max_workers == 0:
max_workers = min(cpu_count(), 8) # Default to 8 workers
log("[STREAMING] " + "=" * 60)
log(f"[STREAMING] Starting streaming mode risk table mapping")
log(f"[STREAMING] Job type: {job_type}")
log(f"[STREAMING] Max workers: {max_workers}")
log(f"[STREAMING] Categorical fields: {cat_field_list}")
log("[STREAMING] " + "=" * 60)
# ========================================================================
# PASS 1: Collect Risk Table Statistics (Training Only)
# ========================================================================
if job_type == "training":
log("[STREAMING] PASS 1: Collecting risk table statistics from train split...")
train_shards = find_split_shards(input_dir, "train", log)
risk_tables = collect_risk_table_statistics_pass1(
train_shards,
signature_columns,
cat_field_list,
label_name,
smooth_factor,
count_threshold,
max_unique_threshold,
log,
)
# Save risk table artifacts
if model_artifacts_output_dir:
artifacts_path = Path(model_artifacts_output_dir)
artifacts_path.mkdir(parents=True, exist_ok=True)
# Save risk_table_map.pkl
risk_table_path = artifacts_path / RISK_TABLE_FILENAME
with open(risk_table_path, "wb") as f:
pkl.dump(risk_tables, f)
log(f"[STREAMING] Saved risk table to {risk_table_path}")
# Save hyperparameters
hyperparams_path = artifacts_path / HYPERPARAMS_FILENAME
with open(hyperparams_path, "w") as f:
json.dump(hyperparams, f, indent=2)
log(f"[STREAMING] Saved hyperparameters to {hyperparams_path}")
else:
# Non-training: Load risk tables
if not model_artifacts_input_dir:
raise ValueError(f"model_artifacts_input_dir required for {job_type} mode")
risk_table_path = Path(model_artifacts_input_dir) / RISK_TABLE_FILENAME
if not risk_table_path.exists():
raise FileNotFoundError(f"Risk table not found: {risk_table_path}")
log(f"[STREAMING] Loading risk tables from {risk_table_path}")
with open(risk_table_path, "rb") as f:
risk_tables = pkl.load(f)
log(f"[STREAMING] Loaded {len(risk_tables)} risk tables")
# ========================================================================
# PASS 2: Process Each Split Independently in Parallel
# ========================================================================
log("[STREAMING] PASS 2: Processing splits in parallel...")
# Determine which splits to process
if job_type == "training":
splits_to_process = ["train", "val", "test"]
else:
# Single split mode (validation, testing, calibration)
splits_to_process = [job_type]
total_stats = {}
for split_name in splits_to_process:
log(f"[STREAMING] Processing {split_name} split...")
# Find shards for this split
split_shards = find_split_shards(input_dir, split_name, log)
# Auto-detect format from first shard (mirrors batch mode behavior)
output_format = detect_shard_format(split_shards[0])
log(f"[STREAMING] Detected format: {output_format}")
# Build global context for this split
global_context = {
"split_name": split_name,
"risk_tables": risk_tables,
}
# Prepare arguments for parallel processing
shard_args = [
(
shard,
extract_shard_number(shard),
global_context,
output_path,
signature_columns,
output_format,
)
for shard in split_shards
]
# Process shards in parallel using multiprocessing Pool
log(
f"[STREAMING] Processing {len(shard_args)} shards from {split_name} "
f"with {max_workers} workers"
)
with Pool(processes=max_workers) as pool:
results = pool.map(process_shard_end_to_end_risk_mapping, shard_args)
# Aggregate results for this split
split_total = sum(r.get(split_name, 0) for r in results)
total_stats[split_name] = split_total
log(f"[STREAMING] Completed {split_name} split: {split_total:,} rows")
log("[STREAMING] " + "=" * 60)
log(f"[STREAMING] Complete! Row distribution: {total_stats}")
log("[STREAMING] " + "=" * 60)
return total_stats
[docs]
def find_split_shards(
input_dir: str, split_name: str, log_func: Callable
) -> List[Path]:
"""
Find all input shards in a specific split subdirectory.
Used when input data is organized as:
input_dir/
train/part-00000.csv, part-00001.csv, ...
val/part-00000.csv, part-00001.csv, ...
test/part-00000.csv, part-00001.csv, ...
Args:
input_dir: Base input directory
split_name: Split subdirectory name ("train", "val", "test", etc.)
log_func: Logging function
Returns:
Sorted list of shard paths from the split subdirectory
Raises:
RuntimeError: If split subdirectory or shards not found
"""
split_dir = Path(input_dir) / split_name
if not split_dir.exists():
raise RuntimeError(f"Split subdirectory not found: {split_dir}")
patterns = [
"part-*.csv",
"part-*.csv.gz",
"part-*.json",
"part-*.json.gz",
"part-*.parquet",
"part-*.snappy.parquet",
"part-*.parquet.gz",
]
all_shards = sorted([p for pat in patterns for p in split_dir.glob(pat)])
if not all_shards:
raise RuntimeError(f"No shards found in {split_dir}")
log_func(f"[STREAMING] Found {len(all_shards)} shards in {split_name} split")
return all_shards
[docs]
def write_shard_file(df: pd.DataFrame, output_path: Path, output_format: str) -> None:
"""
Write a DataFrame to a shard file in the specified format.
Creates parent directories if needed.
Args:
df: DataFrame to write
output_path: Full path for output file (including filename)
output_format: Format to write ('csv', 'tsv', or 'parquet')
Raises:
ValueError: If output_format is not supported
"""
# Create parent directory if needed
output_path.parent.mkdir(parents=True, exist_ok=True)
if output_format == "csv":
df.to_csv(output_path, index=False)
elif output_format == "tsv":
df.to_csv(output_path, sep="\t", index=False)
elif output_format == "parquet":
df.to_parquet(output_path, index=False)
else:
raise ValueError(f"Unsupported output format: {output_format}")
def _read_file_to_df(
file_path: Path, column_names: Optional[List[str]] = None
) -> pd.DataFrame:
"""
Read a single file (CSV, TSV, JSON, Parquet) into a DataFrame.
Simplified version for streaming mode - handles common formats.
Args:
file_path: Path to file
column_names: Optional column names (for CSV/TSV files)
Returns:
DataFrame from file
"""
suffix = file_path.suffix.lower()
if suffix == ".csv" or (suffix == ".gz" and file_path.stem.endswith(".csv")):
if column_names:
return pd.read_csv(file_path, names=column_names, header=0)
return pd.read_csv(file_path)
elif suffix == ".tsv" or (suffix == ".gz" and file_path.stem.endswith(".tsv")):
if column_names:
return pd.read_csv(file_path, sep="\t", names=column_names, header=0)
return pd.read_csv(file_path, sep="\t")
elif suffix == ".parquet" or suffix.endswith(".parquet"):
return pd.read_parquet(file_path)
else:
# Default to CSV
return pd.read_csv(file_path)
# --- File I/O Helper Functions with Format Preservation ---
def _detect_file_format(split_dir: Path, split_name: str) -> tuple:
"""
Detect the format of processed data file.
Returns:
Tuple of (file_path, format) where format is 'csv', 'tsv', or 'parquet'
"""
# Try different formats in order of preference
formats = [
(f"{split_name}_processed_data.csv", "csv"),
(f"{split_name}_processed_data.tsv", "tsv"),
(f"{split_name}_processed_data.parquet", "parquet"),
]
for filename, fmt in formats:
file_path = split_dir / filename
if file_path.exists():
return file_path, fmt
raise RuntimeError(
f"No processed data file found in {split_dir}. "
f"Looked for: {[f[0] for f in formats]}"
)
[docs]
def load_json_config(config_path: str) -> Dict[str, Any]:
"""
Loads a JSON configuration file.
Args:
config_path: Path to the JSON configuration file
Returns:
Dict containing the loaded JSON configuration
Raises:
FileNotFoundError: If the configuration file doesn't exist
PermissionError: If the configuration file can't be accessed due to permissions
json.JSONDecodeError: If the configuration file contains invalid JSON
Exception: For other unexpected errors
"""
try:
with open(config_path, "r") as file:
return json.load(file)
except FileNotFoundError as e:
logger.error(f"Configuration file not found at {config_path}: {str(e)}")
raise
except PermissionError as e:
logger.error(
f"Permission denied when accessing configuration at {config_path}: {str(e)}"
)
raise
except json.JSONDecodeError as e:
logger.error(f"Invalid JSON format in configuration at {config_path}: {str(e)}")
raise
except Exception as e:
logger.error(
f"Unexpected error loading configuration from {config_path}: {str(e)}"
)
raise
[docs]
def validate_categorical_fields(
df: pd.DataFrame, cat_field_list: List[str], max_unique_threshold: int = 100
) -> List[str]:
"""
Validate that fields in cat_field_list are suitable for risk mapping.
Args:
df: DataFrame containing the data
cat_field_list: List of categorical field names
max_unique_threshold: Maximum unique values threshold for suitability (default: 100)
Returns:
List of valid fields for risk mapping
"""
valid_fields = []
for field in cat_field_list:
if field not in df.columns:
logger.warning(f"Field '{field}' not found in data, skipping")
continue
# Check if field is categorical or has few unique values
unique_count = df[field].nunique()
if (
pd.api.types.is_categorical_dtype(df[field])
or unique_count < max_unique_threshold
):
valid_fields.append(field)
logger.info(
f"Field '{field}' is valid for risk mapping ({unique_count} unique values)"
)
else:
logger.warning(
f"Field '{field}' may not be suitable for risk mapping ({unique_count} unique values, threshold={max_unique_threshold})"
)
return valid_fields
[docs]
class OfflineBinning:
"""
A class to create risk tables for categorical features.
Risk tables map categorical values to numerical risk scores based on
their correlation with the target variable.
"""
def __init__(self, cat_field_list: List[str], target_field: str):
self.risk_tables = {}
self.target = target_field
self.variables = cat_field_list
[docs]
def fit(
self, df: pd.DataFrame, smooth_factor: float = 0, count_threshold: int = 0
) -> None:
"""Fits the risk tables based on the provided dataframe."""
# Drop any -1 or NaN target rows for fitting
fit_df = df.loc[(df[self.target] != -1) & (~df[self.target].isnull())].copy()
default_risk = float(fit_df[self.target].mean())
smooth_samples = int(len(fit_df) * smooth_factor)
for var in self.variables:
if var not in fit_df.columns:
continue
# Infer data type from pandas
if pd.api.types.is_numeric_dtype(fit_df[var]):
data_mode = "numeric"
else:
data_mode = "categorical"
self.risk_tables[var] = {
"varName": var,
"type": "categorical",
"mode": data_mode,
"default_bin": default_risk,
}
if fit_df[var].isnull().all():
self.risk_tables[var]["bins"] = {}
continue
risk_table = self._create_risk_table(
fit_df, var, default_risk, smooth_samples, count_threshold
)
self.risk_tables[var]["bins"] = risk_table
def _create_risk_table(
self,
df: pd.DataFrame,
variable: str,
default_risk: float,
samples: int,
count_threshold: int,
) -> Dict:
"""Helper to calculate the risk table for a single variable."""
cross_tab = pd.crosstab(
df[variable],
df[self.target].astype(object),
margins=True,
margins_name="_count_",
dropna=False,
)
cross_tab["risk"] = cross_tab.apply(
lambda x: x.get(1, 0.0) / (x.get(1, 0) + x.get(0, 0)), axis=1
)
cross_tab["smooth_risk"] = cross_tab.apply(
lambda x: (
(
(x["_count_"] * x["risk"] + samples * default_risk)
/ (x["_count_"] + samples)
)
if x["_count_"] >= count_threshold
else default_risk
),
axis=1,
)
cross_tab = cross_tab.loc[cross_tab.index != "_count_"]
return dict(zip(cross_tab.index, cross_tab["smooth_risk"]))
[docs]
def load_risk_tables(self, risk_tables: Dict) -> None:
"""Loads pre-existing risk tables."""
self.risk_tables = risk_tables
logger.info(f"Loaded {len(risk_tables)} risk tables")
[docs]
def load_split_data(job_type: str, input_dir: str) -> Dict[str, pd.DataFrame]:
"""
Load data according to job_type with automatic format detection.
For 'training': Loads data from train, test, and val subdirectories
For others: Loads single job_type split
Returns:
Dictionary with DataFrames and detected format stored in '_format' key
"""
input_path = Path(input_dir)
result = {}
if job_type == "training":
# For training, we expect data in train/test/val subdirectories
splits = ["train", "test", "val"]
detected_format = None
for split_name in splits:
split_dir = input_path / split_name
file_path, fmt = _detect_file_format(split_dir, split_name)
# Store format from first split (they should all match)
if detected_format is None:
detected_format = fmt
# Read based on format
if fmt == "csv":
df = pd.read_csv(file_path)
elif fmt == "tsv":
df = pd.read_csv(file_path, sep="\t")
elif fmt == "parquet":
df = pd.read_parquet(file_path)
else:
raise RuntimeError(f"Unsupported format: {fmt}")
result[split_name] = df
result["_format"] = detected_format # Store detected format
logger.info(
f"Loaded training data splits (format={detected_format}): "
f"train={result['train'].shape}, test={result['test'].shape}, val={result['val'].shape}"
)
else:
# For other job types, we expect data in a single directory named after job_type
split_dir = input_path / job_type
file_path, detected_format = _detect_file_format(split_dir, job_type)
# Read based on format
if detected_format == "csv":
df = pd.read_csv(file_path)
elif detected_format == "tsv":
df = pd.read_csv(file_path, sep="\t")
elif detected_format == "parquet":
df = pd.read_parquet(file_path)
else:
raise RuntimeError(f"Unsupported format: {detected_format}")
result[job_type] = df
result["_format"] = detected_format # Store detected format
logger.info(f"Loaded {job_type} data (format={detected_format}): {df.shape}")
return result
[docs]
def save_output_data(
job_type: str, output_dir: str, data_dict: Dict[str, pd.DataFrame]
) -> None:
"""
Save processed data according to job_type, preserving input format.
For 'training': Saves data to train, test, and val subdirectories
For others: Saves to single job_type directory
"""
output_path = Path(output_dir)
# Extract format from data_dict (stored during load)
output_format = data_dict.get("_format", "csv") # Default to CSV if not found
for split_name, df in data_dict.items():
# Skip the format metadata key
if split_name == "_format":
continue
split_output_dir = output_path / split_name
split_output_dir.mkdir(exist_ok=True, parents=True)
# Save in detected format
if output_format == "csv":
output_file = split_output_dir / f"{split_name}_processed_data.csv"
df.to_csv(output_file, index=False)
elif output_format == "tsv":
output_file = split_output_dir / f"{split_name}_processed_data.tsv"
df.to_csv(output_file, sep="\t", index=False)
elif output_format == "parquet":
output_file = split_output_dir / f"{split_name}_processed_data.parquet"
df.to_parquet(output_file, index=False)
else:
raise RuntimeError(f"Unsupported output format: {output_format}")
logger.info(
f"Saved {split_name} data to {output_file} (format={output_format}), shape: {df.shape}"
)
[docs]
def process_data(
data_dict: Dict[str, pd.DataFrame],
cat_field_list: List[str],
label_name: str,
job_type: str,
risk_tables_dict: Optional[Dict] = None,
smooth_factor: float = 0.01,
count_threshold: int = 5,
max_unique_threshold: int = 100,
) -> Tuple[Dict[str, pd.DataFrame], OfflineBinning]:
"""
Core data processing logic for risk table mapping.
Args:
data_dict: Dictionary of dataframes keyed by split name
cat_field_list: List of categorical field names
label_name: Target column name
job_type: Type of job (training, validation, testing, calibration)
risk_tables_dict: Pre-existing risk tables (for non-training jobs)
smooth_factor: Smoothing factor for risk tables (default: 0.01)
count_threshold: Minimum count threshold (default: 5)
max_unique_threshold: Maximum unique values threshold (default: 100)
Returns:
Tuple containing:
- Dictionary of transformed dataframes
- OfflineBinning instance with fitted risk tables
"""
# For training mode, we need to validate categorical fields
if job_type == "training":
# Validate categorical fields on training data
valid_cat_fields = validate_categorical_fields(
data_dict["train"], cat_field_list, max_unique_threshold
)
if not valid_cat_fields:
logger.warning(
"No valid categorical fields found for risk mapping. Using original data."
)
transformed_data = data_dict
# Create empty binner with no fields for artifact consistency
binner = OfflineBinning([], label_name)
logger.info(
"Creating empty risk tables since no valid categorical fields found"
)
else:
# Initialize transformers with valid fields
binner = OfflineBinning(valid_cat_fields, label_name)
logger.info(
"Running in 'training' mode: fitting on train data, transforming all splits"
)
# Fit risk tables on training data only
binner.fit(
data_dict["train"],
smooth_factor=smooth_factor,
count_threshold=count_threshold,
)
# Transform all splits
transformed_data = {}
for split_name, df in data_dict.items():
# Skip the format metadata key
if split_name == "_format":
transformed_data[split_name] = df # Preserve the format key
continue
df_transformed = binner.transform(df)
transformed_data[split_name] = df_transformed
logger.info(
f"Transformed {split_name} data, shape: {df_transformed.shape}"
)
# --- Logic for Non-Training modes (validation, testing, calibration) ---
else:
if not risk_tables_dict:
raise ValueError(
"For non-training job types, risk_tables_dict must be provided"
)
# Create binner with categorical fields from loaded risk tables
cat_fields = [
rt_info.get("varName")
for rt_info in risk_tables_dict.values()
if "varName" in rt_info
]
binner = OfflineBinning(cat_fields, label_name)
binner.load_risk_tables(risk_tables_dict)
logger.info(f"Using pre-trained risk tables with {len(cat_fields)} fields")
# Transform the data
transformed_data = {}
for split_name, df in data_dict.items():
# Skip the format metadata key
if split_name == "_format":
transformed_data[split_name] = df # Preserve the format key
continue
df_transformed = binner.transform(df)
transformed_data[split_name] = df_transformed
logger.info(f"Transformed {split_name} data, shape: {df_transformed.shape}")
return transformed_data, binner
[docs]
def save_artifacts(
binner: OfflineBinning, hyperparams: Dict[str, Any], output_path: Path
) -> None:
"""
Save risk table artifacts to the specified output path.
Output format matches XGBoost training's risk_table_map.pkl format:
A dictionary mapping variable names to their risk table dictionaries.
Each risk table dict contains 'bins' and 'default_bin' keys.
Args:
binner: OfflineBinning instance with fitted risk tables
hyperparams: Hyperparameter dictionary
output_path: Path to save artifacts to
"""
# Save risk tables in XGBoost training format - using consistent filename
# Format: {variable_name: {"bins": {...}, "default_bin": value, ...}}
bin_output_path = output_path / RISK_TABLE_FILENAME
with open(bin_output_path, "wb") as f:
pkl.dump(binner.risk_tables, f)
logger.info(f"Saved risk table mapping to {bin_output_path}")
logger.info(f"Format: {list(binner.risk_tables.keys())} variables")
logger.info(f"This file can be used as input for non-training jobs")
# Save a copy of hyperparameters with the model
hyperparams_output_path = output_path / HYPERPARAMS_FILENAME
with open(hyperparams_output_path, "w") as f:
json.dump(hyperparams, f, indent=2)
logger.info(f"Saved hyperparameters to {hyperparams_output_path}")
[docs]
def copy_existing_artifacts(src_dir: str, dst_dir: str) -> None:
"""
Copy all existing model artifacts from previous processing steps.
This enables the parameter accumulator pattern where each step:
1. Copies artifacts from previous steps
2. Adds its own artifacts
3. Passes all artifacts to the next step
Args:
src_dir: Source directory containing existing artifacts
dst_dir: Destination directory to copy artifacts to
"""
if not src_dir or not os.path.exists(src_dir):
logger.info(f"No existing artifacts to copy from {src_dir}")
return
os.makedirs(dst_dir, exist_ok=True)
copied_count = 0
for filename in os.listdir(src_dir):
src_file = os.path.join(src_dir, filename)
dst_file = os.path.join(dst_dir, filename)
if os.path.isfile(src_file):
shutil.copy2(src_file, dst_file)
copied_count += 1
logger.info(f" Copied existing artifact: {filename}")
logger.info(f"✓ Copied {copied_count} existing artifact(s) to {dst_dir}")
[docs]
def load_risk_tables(risk_table_path: Path) -> Dict:
"""
Load risk tables from a pickle file.
Args:
risk_table_path: Path to the risk tables file
Returns:
Dictionary of risk tables
"""
if not risk_table_path.exists():
raise FileNotFoundError(f"Risk table file not found: {risk_table_path}")
logger.info(f"Loading risk tables from {risk_table_path}")
with open(risk_table_path, "rb") as f:
risk_tables = pkl.load(f)
logger.info(f"Successfully loaded risk tables with {len(risk_tables)} mappings")
return risk_tables
[docs]
def internal_main(
job_type: str,
input_dir: str,
output_dir: str,
hyperparams: Dict[str, Any],
environ_vars: Optional[Dict[str, str]] = None,
model_artifacts_input_dir: Optional[str] = None,
model_artifacts_output_dir: Optional[str] = None,
load_data_func: Callable = load_split_data,
save_data_func: Callable = save_output_data,
) -> Tuple[Dict[str, pd.DataFrame], OfflineBinning]:
"""
Main logic for risk table mapping with dual-mode support (batch/streaming).
Args:
job_type: Type of job (training, validation, testing, calibration)
input_dir: Input directory for data
output_dir: Output directory for processed data
hyperparams: Hyperparameters dictionary loaded from hyperparameters.json
environ_vars: Environment variables dictionary
model_artifacts_input_dir: Directory containing model artifacts from previous steps
model_artifacts_output_dir: Directory to save model artifacts for next steps
load_data_func: Function to load data (for dependency injection in tests)
save_data_func: Function to save data (for dependency injection in tests)
Returns:
Tuple containing:
- Dictionary of transformed dataframes
- OfflineBinning instance with fitted risk tables
"""
# environ_vars is optional: when omitted (e.g. a local/test call), fall back to hyperparameter
# defaults. Normalizing to {} keeps the env-priority `.get()` reads below working.
if environ_vars is None:
environ_vars = {}
output_path = Path(output_dir)
output_path.mkdir(parents=True, exist_ok=True)
# Extract required hyperparameters with defaults
# Environment variables take priority over hyperparameters
cat_field_list = hyperparams.get("cat_field_list", [])
label_name = hyperparams.get("label_name", "target")
smooth_factor = float(
environ_vars.get("SMOOTH_FACTOR", hyperparams.get("smooth_factor", 0.01))
)
count_threshold = int(
environ_vars.get("COUNT_THRESHOLD", hyperparams.get("count_threshold", 5))
)
max_unique_threshold = int(
environ_vars.get(
"MAX_UNIQUE_THRESHOLD", hyperparams.get("max_unique_threshold", 100)
)
)
# Extract streaming mode configuration
enable_true_streaming = (
environ_vars.get("ENABLE_TRUE_STREAMING", "false").lower() == "true"
)
max_workers_str = environ_vars.get("MAX_WORKERS", "0")
max_workers = int(max_workers_str) if max_workers_str else 0
logger.info(
f"Using parameters: label_name={label_name}, "
+ f"smooth_factor={smooth_factor}, count_threshold={count_threshold}, "
+ f"max_unique_threshold={max_unique_threshold}"
)
logger.info(f"Categorical fields from hyperparameters: {cat_field_list}")
logger.info(f"Streaming mode: {enable_true_streaming}, max_workers: {max_workers}")
# Determine model artifacts output directory
artifacts_output_dir = (
Path(model_artifacts_output_dir)
if model_artifacts_output_dir
else output_path / "model_artifacts"
)
artifacts_output_dir.mkdir(parents=True, exist_ok=True)
# Copy existing artifacts from previous steps (parameter accumulator pattern)
if model_artifacts_input_dir:
copy_existing_artifacts(model_artifacts_input_dir, str(artifacts_output_dir))
# ========================================================================
# STREAMING MODE
# ========================================================================
if enable_true_streaming:
logger.info("=" * 60)
logger.info("STREAMING MODE ENABLED")
logger.info("=" * 60)
# Call streaming mode orchestration
stats = process_streaming_mode_risk_mapping(
input_dir=input_dir,
output_dir=output_dir,
signature_columns=None, # Files have headers
job_type=job_type,
hyperparams=hyperparams,
environ_vars=environ_vars,
max_workers=max_workers,
model_artifacts_input_dir=model_artifacts_input_dir,
model_artifacts_output_dir=str(artifacts_output_dir),
log_func=logger.info,
)
logger.info(f"Streaming mode complete! Final statistics: {stats}")
# Return empty data dict and None binner (data written to disk)
return {}, None
# ========================================================================
# BATCH MODE (DEFAULT) - Existing code unchanged
# ========================================================================
logger.info("=" * 60)
logger.info("BATCH MODE ENABLED")
logger.info("=" * 60)
# Load data according to job type
data_dict = load_data_func(job_type, input_dir)
# Load risk tables if needed (non-training modes)
risk_tables_dict = None
if job_type != "training" and model_artifacts_input_dir:
# Use the consistent filename for loading risk tables
risk_table_path = Path(model_artifacts_input_dir) / RISK_TABLE_FILENAME
if risk_table_path.exists():
risk_tables_dict = load_risk_tables(risk_table_path)
logger.info(f"Loaded pre-trained risk tables from {risk_table_path}")
else:
logger.warning(f"Risk tables not found at {risk_table_path}")
# Process the data
transformed_data, binner = process_data(
data_dict=data_dict,
cat_field_list=cat_field_list,
label_name=label_name,
job_type=job_type,
risk_tables_dict=risk_tables_dict,
smooth_factor=smooth_factor,
count_threshold=count_threshold,
max_unique_threshold=max_unique_threshold,
)
# Save processed data
save_data_func(job_type, output_dir, transformed_data)
# Save fitted artifacts
save_artifacts(binner, hyperparams, artifacts_output_dir)
logger.info("Risk-table mapping complete.")
return transformed_data, binner
[docs]
def main(
input_paths: Dict[str, str],
output_paths: Dict[str, str],
environ_vars: Dict[str, str],
job_args: Optional[argparse.Namespace] = None,
) -> Tuple[Dict[str, pd.DataFrame], OfflineBinning]:
"""
Standardized main entry point for risk table mapping script.
Args:
input_paths: Dictionary of input paths with logical names
- "data_input": Input data directory
- "model_artifacts_input": Model artifacts from previous steps (standardized)
- "hyperparameters_s3_uri": Path to hyperparameters (optional)
output_paths: Dictionary of output paths with logical names
- "processed_data": Output directory for processed data
- "model_artifacts_output": Model artifacts output for next steps (standardized)
environ_vars: Dictionary of environment variables
job_args: Command line arguments (optional)
Returns:
Tuple containing:
- Dictionary of transformed dataframes
- OfflineBinning instance with fitted risk tables
"""
try:
# Extract paths from input parameters - required keys must be present
if "input_data" not in input_paths:
raise ValueError("Missing required input path: input_data")
if "processed_data" not in output_paths:
raise ValueError("Missing required output path: processed_data")
# Extract job_type from args
if job_args is None or not hasattr(job_args, "job_type"):
raise ValueError("job_args must contain job_type parameter")
job_type = job_args.job_type
input_dir = input_paths["input_data"]
output_dir = output_paths["processed_data"]
# Get standardized model artifacts paths
model_artifacts_input_dir = input_paths.get("model_artifacts_input")
model_artifacts_output_dir = output_paths.get("model_artifacts_output")
# Log input/output paths for clarity
logger.info(f"Input data directory: {input_dir}")
logger.info(f"Output directory: {output_dir}")
if model_artifacts_input_dir:
logger.info(f"Model artifacts input directory: {model_artifacts_input_dir}")
logger.info(
f"Expected risk table path: {Path(model_artifacts_input_dir) / RISK_TABLE_FILENAME}"
)
if model_artifacts_output_dir:
logger.info(
f"Model artifacts output directory: {model_artifacts_output_dir}"
)
# Load hyperparameters with source directory fallback (same pattern as XGBoost)
# Use provided hyperparameters path, with source directory fallback
if "hyperparameters_s3_uri" in input_paths:
hparam_path = input_paths["hyperparameters_s3_uri"]
if not hparam_path.endswith("hyperparameters.json"):
hparam_path = os.path.join(hparam_path, "hyperparameters.json")
else:
# Fallback to source directory if not provided
hparam_path = "/opt/ml/code/hyperparams/hyperparameters.json"
logger.info(f"Loading hyperparameters from {hparam_path}")
if os.path.exists(hparam_path):
hyperparams = load_json_config(hparam_path)
else:
# FAIL with clear error instead of using arbitrary defaults
raise FileNotFoundError(
f"Hyperparameters file not found at {hparam_path}. "
f"Risk table mapping requires hyperparameters to be provided either via "
f"input channel or in source directory at /opt/ml/code/hyperparams/hyperparameters.json"
)
# Execute the internal main logic
return internal_main(
job_type=job_type,
input_dir=input_dir,
output_dir=output_dir,
hyperparams=hyperparams,
environ_vars=environ_vars,
model_artifacts_input_dir=model_artifacts_input_dir,
model_artifacts_output_dir=model_artifacts_output_dir,
)
except Exception as e:
logger.error(f"Error in risk table mapping: {str(e)}")
logger.error(traceback.format_exc())
raise
if __name__ == "__main__":
try:
parser = argparse.ArgumentParser()
parser.add_argument(
"--job_type",
type=str,
required=True,
choices=["training", "validation", "testing", "calibration"],
help="Type of job to perform",
)
args = parser.parse_args()
# Define standard SageMaker paths based on contract
# Separate data and model artifacts into different subfolders
input_paths = {
"input_data": DEFAULT_INPUT_DIR,
"hyperparameters_s3_uri": DEFAULT_CONFIG_DIR,
}
output_paths = {
"processed_data": DEFAULT_OUTPUT_DIR + "/data",
"model_artifacts_output": DEFAULT_OUTPUT_DIR + "/model_artifacts",
}
# For non-training jobs, add model artifacts input path
if args.job_type != "training":
input_paths["model_artifacts_input"] = DEFAULT_MODEL_ARTIFACTS_DIR
# Environment variables dictionary - read from os.environ
environ_vars = {
"SMOOTH_FACTOR": os.environ.get("SMOOTH_FACTOR", "0.01"),
"COUNT_THRESHOLD": os.environ.get("COUNT_THRESHOLD", "5"),
"MAX_UNIQUE_THRESHOLD": os.environ.get("MAX_UNIQUE_THRESHOLD", "100"),
# Streaming mode configuration
"ENABLE_TRUE_STREAMING": os.environ.get("ENABLE_TRUE_STREAMING", "false"),
"MAX_WORKERS": os.environ.get("MAX_WORKERS", "0"),
}
# Execute the main function with standardized inputs
result, _ = main(input_paths, output_paths, environ_vars, args)
logger.info(f"Risk table mapping completed successfully")
sys.exit(0)
except FileNotFoundError as e:
logger.error(f"File not found error: {str(e)}")
sys.exit(1)
except ValueError as e:
logger.error(f"Value error: {str(e)}")
sys.exit(2)
except Exception as e:
logger.error(f"Error in risk table mapping script: {str(e)}")
logger.error(traceback.format_exc())
sys.exit(3)