After studying this material, you should be able to:
ReducibleMetric contract for batch-wise metric computationConsider a simple Feature Importance (FI) system:
Traditional Pipeline:
─────────────────────────────────────────────────
[Ray Inference] → [Write to S3] → [Read from S3] → [Spark Metrics Job]
↑ ↑ ↑
Overhead 1 Overhead 1 Overhead 2
(Data Write) (Data Read) (Cluster Startup)
─────────────────────────────────────────────────
Two core problems:
| Problem | Description | Cost |
|---|---|---|
| Cluster Startup | Acquiring/warming a new cluster for metrics | Fixed overhead per job |
| Data Movement | Writing predictions to S3, then reading back | Scales with data size |
Key Insight: When running 100+ experiments in parallel (like FI does), these overheads multiply significantly.
Before solving the problem, you must understand why not all metrics can be computed the same way.
┌─────────────────────────────────────────────────────────────┐
│ METRIC TAXONOMY │
├─────────────────┬───────────────────────────────────────────┤
│ DECOMPOSABLE │ Can split across ANY arbitrary batches │
│ (e.g., MAE) │ Rows are fully independent │
├─────────────────┼───────────────────────────────────────────┤
│ HOLISTIC │ Needs global view, but can APPROXIMATE │
│ (e.g., ROC-AUC) │ via fixed-grid aggregation │
├─────────────────┼───────────────────────────────────────────┤
│ GROUP-WISE │ Rows must be colocated by segment │
│ (e.g., P@K) │ No segment may span bucket boundaries │
└─────────────────┴───────────────────────────────────────────┘
Critical Vocabulary:
Pinterest defines a unified interface using an Abstract Base Class:
from abc import ABC, abstractmethod
class ReducibleMetric(ABC):
@abstractmethod
def dataset_preprocessor(self, ds):
"""
PURPOSE: Prepare data layout BEFORE batching
- Decomposable: do nothing (rows independent)
- Holistic: do nothing (rows independent)
- Group-wise: colocate segments into buckets
"""
pass
@abstractmethod
def compute_partials(self, batch):
"""
PURPOSE: Compute intermediate result from ONE batch
- Called once per batch/bucket
- Must return something that can be combined later
"""
pass
@abstractmethod
def reduce_partials(self, partials):
"""
PURPOSE: Combine ALL partial results into final answer
- Called once globally after all batches processed
"""
pass
Why this design?
Batch 1 → compute_partials() → partial_1 ─┐
Batch 2 → compute_partials() → partial_2 ─┤→ reduce_partials() → FINAL METRIC
Batch N → compute_partials() → partial_N ─┘
This enables parallelism: all batches can be processed independently, then merged once.
Formula: MAE = (1/N) × Σ|prediction - label|
Why it's decomposable: Each row's error is independent of all other rows.
Dataset:
┌──────┬──────────────┬──────────────────┬───────┐
│ Row │ model_output │ label_continuous │ error │
├──────┼──────────────┼──────────────────┼───────┤
│ 1 │ 0.80 │ 0.90 │ 0.10 │
│ 2 │ 0.60 │ 0.75 │ 0.15 │
│ 3 │ 0.90 │ 0.60 │ 0.30 │
│ 4 │ 0.70 │ 0.80 │ 0.10 │ ← Batch A: Σ=0.65, n=4
├──────┼──────────────┼──────────────────┼───────┤
│ 5 │ 0.40 │ 0.20 │ 0.20 │
│ 6 │ 0.30 │ 0.10 │ 0.20 │
│ 7 │ 0.20 │ 0.40 │ 0.20 │
│ 8 │ 0.10 │ 0.30 │ 0.20 │ ← Batch B: Σ=0.80, n=4
└──────┴──────────────┴──────────────────┴───────┘
Global MAE = (0.65 + 0.80) / (4 + 4) = 0.18125 ✓
class MAE(ReducibleMetric):
def dataset_preprocessor(self, ds):
return ds # No preparation needed
def compute_partials(self, batch):
error = np.abs(batch["model_output"] - batch["label_continuous"])
return error.sum(), len(error)
# Returns: (sum_of_errors, row_count)
def reduce_partials(self, partials):
error, count = np.sum(partials, axis=0)
return error / count
# Final division happens ONCE globally
Key Insight: Never divide inside compute_partials. Carry forward (sum, count) and divide once in reduce_partials.
Why it's holistic: AUC depends on how rows rank against ALL other rows globally.
The Problem with Naive Batching:
Batch A AUC = 1.00
Batch B AUC = 0.75
Average = 0.875 ← WRONG (not 0.9375)
Why wrong? A row in Batch A might rank below a row in Batch B.
Averaging ignores cross-batch rankings.
Solution: Confusion Matrix Approximation
Instead of sorting globally, use fixed thresholds:
Thresholds: [0.0, 0.25, 0.50, 0.75, 1.0]
For each threshold T, count:
TP = predicted ≥ T AND actually positive
FP = predicted ≥ T AND actually negative
TN = predicted < T AND actually negative
FN = predicted < T AND actually positive
These counts ARE additive across batches!
Batch A confusion matrix (at threshold 0.5):
TP=3, FP=1, TN=0, FN=0
Batch B confusion matrix (at threshold 0.5):
TP=1, FP=0, TN=3, FN=0
Global = Batch A + Batch B:
TP=4, FP=1, TN=3, FN=0
class ROCAUC(ReducibleMetric):
def dataset_preprocessor(self, ds):
return ds # Rows are independent
def compute_partials(self, batch):
pred = batch["model_output"]
truth = batch["label_binarized"].astype(bool)
# Shape: (num_thresholds, num_rows)
positive = pred[None, :] >= self.thresholds[:, None]
# Returns confusion matrix: shape (num_thresholds, 4)
return np.stack([
(positive & truth).sum(1), # TP per threshold
(positive & ~truth).sum(1), # FP per threshold
(~positive & ~truth).sum(1), # TN per threshold
(~positive & truth).sum(1) # FN per threshold
], axis=1)
def reduce_partials(self, partials):
tp, fp, tn, fn = np.sum(partials, axis=0).T
tpr = tp / (tp + fn) # True Positive Rate
fpr = fp / (fp + tn) # False Positive Rate
# Integrate using trapezoid rule (ascending FPR order)
return np.trapz(np.r_[0, tpr[::-1], 1],
np.r_[0, fpr[::-1], 1])
Important Caveat:
More thresholds → Better approximation of true AUC
Step size 0.01 → Often recovers exact value in practice
But: NOT guaranteed to be exact in all cases
Formula: For each segment, rank items by score, take top K, count relevant items, divide by K. Average across segments.
Example with K=2:
Segment ReqA (Batch A):
Ranked: [row3(0.90), row1(0.80), row4(0.70), row2(0.60)]
Top 2: [row3, row1] → labels [1, 1] → P@2 = 2/2 = 1.00
Segment ReqB (Batch B):
Ranked: [row5(0.40), row6(0.30), row7(0.20), row8(0.10)]
Top 2: [row5, row6] → labels [1, 0] → P@2 = 1/2 = 0.50
Final P@2 = (1.00 + 0.50) / 2 = 0.75
The Two Scaling Problems:
Problem 1: Variable segment sizes
Segment A: 4 rows → small, inefficient batch
Segment B: 4 rows → small, inefficient batch
Solution: Pack multiple segments into one bucket
Problem 2: Segments must not be split
If ReqA spans two buckets → wrong ranking within segment
Solution: bucket_complete_segments() ensures no splits
class PrecisionAtK(ReducibleMetric):
def dataset_preprocessor(self, ds):
# CRITICAL: Must happen before Ray batching
return bucket_complete_segments(
ds,
by="request_id", # Group key
)
# Guarantees: all rows for each request_id
# are in the same bucket
def compute_partials(self, bucket):
scores = []
for segment in group_by(bucket, "request_id"):
ranking = np.argsort(-segment["model_output"])
top_k = ranking[:self.k]
relevant = segment["label_binarized"][top_k].sum()
scores.append(relevant / self.k)
return sum(scores), len(scores)
# Returns: (total_score, segment_count)
def reduce_partials(self, partials):
score, count = np.sum(partials, axis=0)
return score / count
UNFUSED (Sequential):
─────────────────────────────────────────────────────────────
Time →
[GPU: Inference B1][GPU: Inference B2]...[GPU: Inference BN]
↓ S3 write/read
[CPU: Metrics B1][CPU: Metrics B2]...
─────────────────────────────────────────────────────────────
Problem: CPU workers idle during inference
Extra S3 read/write cost
Extra cluster startup cost
FUSED (Pipelined):
─────────────────────────────────────────────────────────────
Time →
[GPU: Inference B1] → [CPU: Metrics B1]
[GPU: Inference B2] → [CPU: Metrics B2]
[GPU: Inference B3] → [CPU: Metrics B3]
─────────────────────────────────────────────────────────────
Benefit: Metrics overlap with inference
Shared object store (no S3 round-trip)
One cluster startup instead of two
The ReducibleMetric contract makes fusion possible:
compute_partials() runs per bucket immediately after inferencereduce_partials() runs once after all buckets completeAt E=100 experiments:
┌──────────────┬──────────┬──────────┬──────────────┐
│ Dataset Size │ Unfused │ Fused │ Time Saved │
├──────────────┼──────────┼──────────┼──────────────┤
│ Small ~30K │ baseline │ -25% │ significant │
│ Medium ~300K │ baseline │ -28% │ significant │
│ Large ~3M │ baseline │ -32% │ significant │
└──────────────┴──────────┴──────────┴──────────────┘
Fused runtime ≈ 43 min (fixed) + 41 sec × (num_experiments)
Unfused metrics ≈ 11 sec × (num_experiments) [absorbed by fusion]
At E=300:
Fused: 43 + (41 × 300)/60 ≈ 248 min
Unfused: 43 + (52 × 300)/60 ≈ 303 min
┌─────────────────────────────────────────────────────────┐
│ COMPLETE PICTURE │
│ │
│ 1. Classify metric → determines data layout needs │
│ │
│ 2. dataset_preprocessor() │
│ └─ Group-wise only: bucket_complete_segments() │
│ │
│ 3. Ray distributes buckets to workers │
│ └─ compute_partials() runs per bucket │
│ └─ Happens DURING inference (fused) │
│ │
│ 4. reduce_partials() combines all partials │
│ └─ One final pass, globally │
│ │
│ 5. Result: same accuracy, fraction of the cost │
└─────────────────────────────────────────────────────────┘
The core principle: Treat decomposability as a first principle — design every metric so its computation can be split, parallelized, and merged, enabling fusion with inference and eliminating redundant data movement.