Batch Prediction Pipelines

You’ve trained a model, tuned it, and celebrated its accuracy — but now comes the unglamorous part that separates a demo from a production system: applying that model to thousands or millions of rows of data without melting your machine or losing your sanity. Scripting a for loop over every record is a one-way ticket to slow, fragile, and unmaintainable code. In this lesson, you’ll learn to build batch prediction pipelines that process large datasets efficiently, handle errors gracefully, and scale from a local CSV to a production job — the skills that AI engineers use every day.

The problem this lesson solves

Batch prediction means applying a machine learning model to a large set of records all at once, rather than one at a time. It’s how you score a customer database, label thousands of images, or generate recommendations overnight. The challenge? Your model may be slow, your data may be messy, and your infrastructure may not handle a million rows in a single Python list.

A naive approach — like looping over every row and calling model.predict() — is simple but painfully inefficient. It wastes I/O, ignores vectorization, and crashes when one row has missing or unexpected data. You need a structured way to load, transform, predict, and save in a pipeline that is reusable, testable, and debuggable.

By the end of this lesson, you’ll understand the core concepts behind build batch prediction pipelines and complete a hands-on exercise that turns a scrappy script into a production-ready pipeline.

Core concept / mental model

Think of a batch prediction pipeline as a factory assembly line. Raw materials (your input data) enter on a conveyor belt. Each station does one job: clean the data, transform it into features, feed it to the model, and package the results. The pipeline runs until every box is filled, then the whole batch moves to shipping.

In code terms, this maps to a sequence of stages: load → preprocess → predict → postprocess → save. Each stage is a Python function that takes an input and returns an output, and the pipeline chains them together. This modular design makes it easy to test each piece independently, swap stages, and handle failures without losing the entire batch.

Key definitions to hold in mind:

• Batch: A set of records processed together (e.g., 10,000 rows in a DataFrame).
• Pipeline: A series of processing steps that transform raw inputs into predictions.
• Vectorization: Using array operations (like pandas or NumPy) to apply a function to many rows simultaneously, instead of looping.
• Chunking: Dividing a large dataset into smaller pieces to manage memory and enable progress tracking.
• Idempotency: When running the pipeline twice, you get the same result (important for retries and reproducibility).

Pro tip: A good pipeline is like a recipe — each step is explicit, and you can swap ingredients (e.g., a different model) without rewriting the whole process.

How it works step by step

Building a batch prediction pipeline follows a logical sequence that applies to most ML models. Here’s the blueprint:

1. Load raw data: Read your input from a source like CSV, JSON, or a database. In production, this might be a Parquet file in S3 or a SQL query. For this lesson, we’ll use CSV.
2. Preprocess and transform: Handle missing values, encode categorical variables, scale numerical features, and apply the same transformations you used at training time. This step is crucial — the model expects the same format it saw during training.
3. Predict: Feed the processed batch to your model. This could be a single predict() call on a DataFrame or a loop over chunks if memory is tight. Many libraries (like scikit-learn and XGBoost) support batch predictions natively.
4. Postprocess: Convert raw scores into useful outputs — e.g., class labels, probabilities, or rankings. You might also add business logic like thresholds.
5. Save results: Write predictions to a file or database. Consider including metadata like timestamps and input keys for traceability.

A common pattern is to encapsulate this in a function that takes a file path and returns a result path, allowing you to reuse it across different datasets. Each step should be independent — if the prediction step fails, you can restart only that step, not reload everything.

To handle large datasets, you’ll often use chunking: read data in chunks (e.g., 10,000 rows at a time), process each chunk, and append results. This keeps memory usage flat, even with millions of records.

Hands-on walkthrough

Let’s build a simple but complete batch prediction pipeline using scikit-learn. We’ll generate synthetic data, train a quick model, and then process a large CSV file in chunks. Make sure you have pandas, numpy, and scikit-learn installed.

Step 1: Set up your environment

pip install pandas numpy scikit-learn

Step 2: Create synthetic training data and train a model

import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestClassifier

# Generate synthetic data: 1000 samples, 5 features, binary target
np.random.seed(42)
X = np.random.rand(1000, 5)
y = (X[:, 0] + 2 * X[:, 1] > 1.5).astype(int)

df = pd.DataFrame(X, columns=[f"feature_{i}" for i in range(5)])
df["target"] = y

# Save training data
df.to_csv("training_data.csv", index=False)

# Train model
X_train = df.drop("target", axis=1)
y_train = df["target"]
model = RandomForestClassifier(n_estimators=100, random_state=42)
model.fit(X_train, y_train)
print("Model trained.")

Output:

Model trained.

Step 3: Build the batch prediction pipeline

Now we’ll create a pipeline that processes new data in chunks, logs progress, and handles errors gracefully.

import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestClassifier
import pickle

def predict_batch(input_path: str, output_path: str, chunk_size: int = 1000):
    """Run batch predictions on a CSV file, writing results to output_path."""
    # Load the model once
    with open("model.pkl", "rb") as f:
        model = pickle.load(f)

    # Prepare output file (write header)
    header_written = False

    # Iterate over chunks
    for chunk in pd.read_csv(input_path, chunksize=chunk_size):
        # Assume the CSV has the same feature columns (no target)
        features = chunk.drop(columns=["id"], errors="ignore")  # drop id if present
        predictions = model.predict(features)
        probabilities = model.predict_proba(features)[:, 1]

        result_df = chunk.copy()
        result_df["prediction"] = predictions
        result_df["probability"] = probabilities

        # Write to output, append after first chunk
        result_df.to_csv(output_path, mode="a", header=not header_written, index=False)
        header_written = True

        print(f"Processed {len(chunk)} rows")

    print("Batch prediction complete.")

# Save the model after training
with open("model.pkl", "wb") as f:
    pickle.dump(model, f)

# Create a large dataset to predict on (e.g., 10,000 rows)
new_data = pd.DataFrame(np.random.rand(10000, 5), columns=[f"feature_{i}" for i in range(5)])
new_data.insert(0, "id", range(10000))
new_data.to_csv("new_data.csv", index=False)

# Run the pipeline
predict_batch("new_data.csv", "predictions.csv", chunk_size=2000)

Expected output (truncated):

Processed 2000 rows
Processed 2000 rows
Processed 2000 rows
Processed 2000 rows
Processed 2000 rows
Batch prediction complete.

Step 4: Add error handling

Real-world data is messy. Let’s enhance the pipeline to skip rows that cause errors and log them instead of crashing.

import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestClassifier
import pickle

def predict_batch_safe(input_path: str, output_path: str, error_log: str, chunk_size: int = 1000):
    """Run batch predictions with error handling."""
    with open("model.pkl", "rb") as f:
        model = pickle.load(f)

    header_written = False
    error_count = 0

    for chunk in pd.read_csv(input_path, chunksize=chunk_size):
        features = chunk.drop(columns=["id"], errors="ignore")
        try:
            predictions = model.predict(features)
            probabilities = model.predict_proba(features)[:, 1]
        except Exception as e:
            # Log errors and skip this chunk
            with open(error_log, "a") as err_f:
                err_f.write(f"Chunk failed: {str(e)}\n")
            error_count += 1
            continue

        result_df = chunk.copy()
        result_df["prediction"] = predictions
        result_df["probability"] = probabilities
        result_df.to_csv(output_path, mode="a", header=not header_written, index=False)
        header_written = True
        print(f"Processed chunk with {len(chunk)} rows")

    print(f"Pipeline finished with {error_count} chunk(s) failed. See {error_log}.")

# Usage (assuming model.pkl exists from previous step)
predict_batch_safe("new_data.csv", "predictions_safe.csv", "errors.log", chunk_size=2000)

In this version, if a chunk contains invalid data that causes an exception, the pipeline logs the error and continues with the next chunk — crucial for production where you can’t afford a single bad row to kill the job.

Compare options / when to choose what

There isn’t one “right” way to build a prediction pipeline; it depends on your scale and environment. Here’s a comparison of common approaches:

For most Python AI engineering, chunking is the sweet spot: it’s simple, portable, and handles data larger than RAM. If you need scale-out, consider parallel libraries like Dask or Polars (string).

Pro tip: When choosing a chunk size, aim for a size that balances memory and I/O — typically 10,000–50,000 rows for CSV files, but test on your machine.

Troubleshooting & edge cases

Even a well-designed pipeline can hit snags. Here are common problems and how to fix them:

• MemoryError when reading large CSV: You’re loading the entire file at once. Solution: Use pd.read_csv(..., chunksize=...) to process in chunks, or use dtype to specify column types and reduce memory. python for chunk in pd.read_csv("huge.csv", chunksize=10000, low_memory=False): # process
• Feature mismatch between training and prediction: The model expects columns feature_0 to feature_4, but your new data has different names or column order. Always apply the same preprocessing pipeline (e.g., using the same ColumnTransformer) and check column order.
• Missing values in new data: The model may fail if it sees NaN. Impute missing values (mean/median) using the same statistics from training — better yet, use a Pipeline with an imputer to keep consistent.
• Slow predictions due to Python overhead: If you’re calling model.predict() on each row, that’s slow. Switch to batch predictions on arrays or use predict on the whole DataFrame. For custom models, use NumPy vectorization.
• File locking issues when writing: If multiple processes write to the same output file, they may conflict. Use separate output files per worker or a lock mechanism.
• Reproducibility: If you set random_state at training, but your pipeline uses shuffling elsewhere, set a global seed.

Example of fixing a feature mismatch:

# Check if all expected columns are present
required_cols = [f"feature_{i}" for i in range(5)]
if not set(required_cols).issubset(new_data.columns):
    raise ValueError(f"Missing columns: {set(required_cols) - set(new_data.columns)}")
new_data = new_data[required_cols]  # ensure correct order

What you learned & what's next

You now understand the core idea behind build batch prediction pipelines: create modular, chunked, error-tolerant processes that turn raw data into predictions. You learned how to implement a pipeline in Python, compare approaches, and troubleshoot common issues. You completed a practical exercise where you trained a model, processed a large CSV in chunks, and added error handling.

Next in this track, you’ll build on these skills by adding model monitoring or retraining workflows — but for now, you’ve got a solid pipeline that can handle real-world data sizes. Keep experimenting with different chunk sizes, models, and data sources.

Final thought: A batch prediction pipeline isn’t just about predicting — it’s about building trust that your model can run reliably, at scale, without surprises.