Scaling ML Model Evaluation with Ray: Metrics at Scale

After studying this material, you should be able to:

  1. Explain why model evaluation becomes an infrastructure challenge at scale
  2. Classify metrics into decomposable, holistic, and group-wise categories
  3. Implement the ReducibleMetric contract for batch-wise metric computation
  4. Describe how fusing inference with metric computation reduces end-to-end runtime
  5. Compare unfused vs. fused pipeline architectures and their tradeoffs

Step-by-Step Study Material

Step 1: Understanding the Problem

Why Is Model Evaluation Hard at Scale?

Consider 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:

ProblemDescriptionCost
Cluster StartupAcquiring/warming a new cluster for metricsFixed overhead per job
Data MovementWriting predictions to S3, then reading backScales with data size

Key Insight: When running 100+ experiments in parallel (like FI does), these overheads multiply significantly.


Step 2: The Three Metric Categories

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:

  • Batch: Arbitrary group of rows (compute unit for one Ray task)
  • Segment: Semantic unit (e.g., one user request)
  • Bucket: Group-preserving batch — complete segments only, no segment crosses boundaries

Step 3: The Metric Contract (Design Pattern)

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.


Step 4: Implementing Each Metric Type

4A: Decomposable Metric — MAE

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.


4B: Holistic Metric — ROC-AUC

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

4C: Group-wise Metric — Precision@K

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

Step 5: Fusing Inference with Metrics

The Architecture Shift

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

Why This Works

The ReducibleMetric contract makes fusion possible:

  1. compute_partials() runs per bucket immediately after inference
  2. Results stay in Ray's shared object store (no S3 write)
  3. reduce_partials() runs once after all buckets complete

Step 6: Results and Tradeoffs

Performance Gains

At E=100 experiments:
┌──────────────┬──────────┬──────────┬──────────────┐
│ Dataset Size │ Unfused  │  Fused   │  Time Saved  │
├──────────────┼──────────┼──────────┼──────────────┤
│ Small ~30K   │ baseline │ -25%     │ significant  │
│ Medium ~300K │ baseline │ -28%     │ significant  │
│ Large ~3M    │ baseline │ -32%     │ significant  │
└──────────────┴──────────┴──────────┴──────────────┘

Linear Scaling

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

Summary: Mental Model

┌─────────────────────────────────────────────────────────┐
│              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.

More to study