From Spark model to production: persistence, MLflow tracking, and batch scoring

LEARN · DISTRIBUTED ML WITH PYSPARK

A trained Spark ML model is only the beginning of a machine learning lifecycle. In production, the hard problems appear after training:

How do you preserve the exact pipeline that created the model?
How do you reproduce a successful experiment months later?
How do you score new data consistently?
How do you prove which parameters produced a prediction?

A reliable workflow connects three pieces:

Spark MLlib for scalable model training and pipeline execution.
Spark ML persistence for saving and loading complete models.
MLflow tracking for experiment history, parameters, metrics, and artifacts.

The result is a repeatable path:

Data → Feature pipeline → Training → Saved model → MLflow tracking → Batch scoring → Predictions

This tutorial uses current Spark ML APIs and MLflow tracking patterns. Spark continues to ship MLlib as part of the Spark project, and current Spark releases include MLlib support. MLflow 3.x is the current major generation, with MLflow 3.0 released in 2025.

Why saving only the model is a production mistake

A common first version of a machine learning workflow looks like this:

Train a model in a notebook.
Keep the model object in memory.
Export a prediction file.
Forget the exact settings used.

This works during exploration but creates operational problems:

Feature transformations may not be identical during inference.
Training parameters may disappear.
A successful experiment may be impossible to reproduce.
Debugging a prediction issue becomes guesswork.

In Spark, the solution is usually to save the entire PipelineModel, not only the final estimator.

A pipeline can contain:

Feature transformers.
Vector assembly.
Encoders.
The trained algorithm.
Additional preprocessing stages.

When the full pipeline is saved, scoring uses the same transformation logic that was used during training.

Training and saving a Spark ML pipeline

The following example creates a simple classification pipeline, trains it, and persists the resulting model.

from pyspark.sql import SparkSession
from pyspark.ml import Pipeline
from pyspark.ml.classification import LogisticRegression
from pyspark.ml.feature import VectorAssembler

spark = SparkSession.builder.appName("model-training").getOrCreate()

training = spark.createDataFrame(
    [
        (0, 1.0, 2.0),
        (1, 2.0, 1.0),
        (0, 0.5, 1.5),
        (1, 3.0, 2.5),
    ],
    ["label", "feature_a", "feature_b"],
)

assembler = VectorAssembler(
    inputCols=["feature_a", "feature_b"],
    outputCol="features",
)

classifier = LogisticRegression(
    featuresCol="features",
    labelCol="label",
    maxIter=20,
)

pipeline = Pipeline(
    stages=[
        assembler,
        classifier,
    ]
)

model = pipeline.fit(training)

model.write().overwrite().save(
    "models/customer_classifier"
)

The saved directory contains Spark metadata and stage information. It is not a Python pickle file and should be handled through Spark ML persistence APIs.

Loading a model for inference

A batch scoring application can reload the complete pipeline:

from pyspark.sql import SparkSession
from pyspark.ml import PipelineModel

spark = SparkSession.builder.appName("model-loading").getOrCreate()

model = PipelineModel.load(
    "models/customer_classifier"
)

new_data = spark.createDataFrame(
    [
        (4.0, 3.0),
        (1.0, 0.5),
    ],
    ["feature_a", "feature_b"],
)

predictions = model.transform(new_data)

predictions.select(
    "feature_a",
    "feature_b",
    "prediction",
    "probability",
).show()

The important detail is that inference does not recreate feature engineering manually. The saved pipeline owns the transformation sequence.

Why MLflow tracking matters

Spark persistence answers one question:

How do I save and reload this model?

MLflow answers a broader production question:

Which model was created, with which configuration, and why should we trust it?

An MLflow run can store:

Hyperparameters.
Metrics.
Tags.
Model artifacts.
Dataset references.
Experiment history.

This metadata becomes essential when teams compare dozens of experiments.

Logging a Spark ML experiment with MLflow

The mlflow.spark integration can log Spark ML models directly.

import mlflow
import mlflow.spark

from pyspark.ml import Pipeline
from pyspark.ml.classification import LogisticRegression
from pyspark.ml.feature import VectorAssembler

assembler = VectorAssembler(
    inputCols=["feature_a", "feature_b"],
    outputCol="features",
)

classifier = LogisticRegression(
    maxIter=30,
    regParam=0.1,
)

pipeline = Pipeline(
    stages=[
        assembler,
        classifier,
    ]
)

with mlflow.start_run() as run:
    mlflow.log_param(
        "maxIter",
        30,
    )

    mlflow.log_param(
        "regParam",
        0.1,
    )

    model = pipeline.fit(training)

    predictions = model.transform(training)

    correct = predictions.filter(
        "prediction = label"
    ).count()

    total = predictions.count()

    accuracy = correct / total

    mlflow.log_metric(
        "training_accuracy",
        accuracy,
    )

    mlflow.spark.log_model(
        model,
        artifact_path="spark-model",
    )

    print("Run ID:", run.info.run_id)

Instead of referring to “the model from last Tuesday’s notebook,” a team can reference a specific MLflow run with recorded parameters and metrics.

The reproducibility rescue: remembering what actually worked

One of the most valuable MLflow features is not the model artifact. It is the small pieces of information humans forget.

A team may remember:

“The older model had stronger regularization.”
“The previous feature set performed better.”
“The validation score was higher.”

Those statements are not enough to rebuild a production model.

A logged MLflow run can preserve:

Regularization values.
Iteration counts.
Feature flags.
Dataset identifiers.
Evaluation metrics.

This creates a reproducibility rescue. When a newer experiment performs worse, engineers can return to the exact configuration that produced the earlier result.

The surprising lesson is that experiment history is often more valuable than the model file itself. A model without context is just an artifact. A model with lineage becomes an engineering asset.

Building a batch scoring job

Production scoring should be automated. A batch scoring job typically:

Loads the approved model.
Reads new data.
Applies the saved transformations.
Writes predictions.

A simple Spark BatchScoring workflow looks like this:

from pyspark.sql import SparkSession
from pyspark.ml import PipelineModel

spark = SparkSession.builder.appName("batch-scoring").getOrCreate()

model_path = "models/customer_classifier"

input_path = "data/incoming/customers"

output_path = "data/predictions/customers"

model = PipelineModel.load(model_path)

customers = spark.read.parquet(
    input_path
)

predictions = model.transform(
    customers
)

result = predictions.select(
    "customer_id",
    "prediction",
    "probability",
)

(
    result.write
    .mode("overwrite")
    .parquet(output_path)
)

Parquet is a practical output format for analytics workloads because it is column-oriented and works well across data platforms.

Real production jobs usually add:

Schema validation.
Data quality checks.
Partitioning rules.
Failure notifications.
Prediction monitoring.

Monitoring after deployment

A saved model can still become inaccurate over time.

Common causes include:

Changing customer behavior.
New data sources.
Pipeline changes.
Feature distribution shifts.

Useful monitoring signals include:

Number of scored records.
Missing feature percentages.
Prediction distribution changes.
Model performance after labels arrive.

MLflow helps compare experiments, but production monitoring requires observing the complete system around the model.

A security lesson from model artifacts

Machine learning artifacts should be treated like software supply-chain components.

A useful reminder came from vulnerabilities in Python serialization ecosystems: loading untrusted serialized objects can become a security risk because some formats can execute code during deserialization.

The practical rules are:

Only load models from trusted storage.
Control who can publish model artifacts.
Keep Spark, MLflow, and dependencies updated.
Use review and promotion workflows before production deployment.

A model file is not just data. It is executable behavior packaged for reuse.

Production checklist

Before moving a Spark MLlib workflow into production, verify:

The complete Spark pipeline is saved.
MLflow records parameters and metrics.
Model artifacts have ownership and version history.
Batch scoring runs automatically.
Input and output schemas are controlled.
Dependencies are pinned and maintained.
Monitoring detects data and prediction changes.

Final thoughts

A production machine learning system is not defined only by the algorithm. Reliable systems combine model persistence, experiment tracking, and automated scoring.

Spark ML provides the scalable execution layer. MLflow provides the history and traceability needed for responsible model development. Together, they create a workflow where models can move from notebooks into dependable production pipelines.

Start your next Spark ML project by saving complete pipelines, logging every important experiment detail, and building batch scoring from the beginning. The effort invested in reproducibility will pay off the first time you need to audit, restore, or improve a model.