Scaling ML Model Evaluation: Fusing Inference and Metrics with Ray and Spark

Evaluating Machine Learning models at scale presents significant infrastructure challenges, especially in systems running numerous experiments in parallel, like feature importance platforms. Traditionally, inference and model evaluation are separate stages, leading to inefficiencies due to data serialization/deserialization and cluster spin-up times. Pinterest's solution optimizes this by fusing these stages.

The Challenge: Unfused Inference and Evaluation

In an unfused pipeline, ML model predictions are first generated (e.g., using Ray) and written to storage (e.g., S3). A separate Apache Spark job then reads this data to compute evaluation metrics. This creates two primary overheads:

• Cluster Startup: The time and resources required to acquire and warm up a new cluster for the metrics computation job.
• Data Movement: The costly process of writing predictions to S3 and then rereading them for metric calculation (serialization/deserialization overhead).

The Solution: Fused Inference and Metrics Compute

Pinterest redesigned its ML infrastructure to fuse inference and metrics computation. This involves calculating metrics per batch directly within the Ray inference job, leveraging Ray's shared object store to avoid cross-stack data handoffs. This strategy reduced end-to-end evaluation time by up to 32% and scaled linearly with the number of experiments.

The ReducibleMetric Contract

A crucial abstraction for this fusion is the `ReducibleMetric` contract. This interface standardizes how different metrics are computed and aggregated, ensuring consistency across execution frameworks like Ray and Spark. It decomposes metric calculation into three abstract methods:

from abc import ABC, abstractmethod

class ReducibleMetric(ABC):
    @abstractmethod
    def dataset_preprocessor(self, ds):
        """Prepare the dataset to satisfy the metric's data layout requirements."""
        pass

    @abstractmethod
    def compute_partials(self, batch):
        """Compute a partial result from one batch"""
        pass

    @abstractmethod
    def reduce_partials(self, partials):
        """Combine partial results into the final result."""
        pass

Handling Different Metric Types

• Decomposable Metrics (e.g., MAE): These metrics can compute partial results from arbitrary batches, and these partials can be simply summed and aggregated globally. No special data arrangement is needed.
• Holistic Metrics (e.g., ROC-AUC): Metrics like ROC-AUC traditionally require a global sort, which is costly. The solution approximates global ROC-AUC using a threshold-based confusion matrix approach, where batch-wise confusion matrices are summed to form a global matrix, from which AUC can be derived.
• Group-wise Metrics (e.g., Precision@K): These metrics depend on data within a semantic unit (e.g., a user request). To handle them efficiently, data is 'bucketed' such that all rows for a given segment are colocated within a single bucket, and no segment spans across buckets. Apache Spark is used for efficient grouping and packing of segments into these buckets, which are then processed by Ray workers.