One Model, Two Worlds — MLflow Models for Real-Time & Batch

The “Works on My Machine” Special

Data scientist trains a model in a Jupyter notebook, gets 97% accuracy, pushes it to GitHub. Then the DevOps team spends two days figuring out what Python version they need, what packages to install, and why the model won’t load in production.

Then the risk team comes along and says “cool, now make it run on our Spark cluster at 2 AM on 10 million rows.”

Now you’ve got two separate deployment targets and zero shared code.

MLflow Models fixes this with flavors.

The Problem

You’re at a fintech company. Your data science team trained a Random Forest to catch fraudulent credit card transactions. Two teams need it:

  • Engineering wants a REST API — real-time blocking of fraudulent transactions at checkout.
  • Risk & Compliance wants a Spark batch job — run every night on millions of historical transactions.

Different runtimes. Different infrastructure. Different teams. Same model.

Without MLflow, you’d rewrite the model code twice — once in FastAPI for the API people, once in PySpark for the batch people.

The MLflow Models Fix

When you save a model with MLflow, it doesn’t just dump a pickle file. It writes a directory structure with an MLmodel file that acts like a Swiss Army knife. It tells every downstream tool:

  • Using Scikit-learn? Here’s the sklearn flavor.
  • Need a generic Python function? Here’s the python_function flavor.
  • On Spark? I got you too.

Same model file. No rewrites.

artifacts/                              # Root directory for all logged model artifacts
├── fraud_detector/                     # Subdirectory holding the fraud detection model artifacts
│   ├── MLmodel                         # Metadata file: defines flavors, signature, dependencies, and environment
│   ├── conda.yaml                      # Conda environment spec pinning every dependency for reproducible serving
│   ├── requirements.txt                # pip requirements file with exact versions used during model training
│   ├── python_env.yaml                 # Python environment descriptor (alternative to conda.yaml for non-conda users)
│   ├── model.pkl                       # Serialized RandomForestClassifier model binary (pickle format)
│   └── input_example.json              # Saved example input row for smoke-testing the deployed endpoint

Step 1: Train and Log

Train your model and log it with a signature (schema) and an input example. This one call generates everything — the flavors, the dependency files, the metadata:

import mlflow.sklearn                   # MLflow's scikit-learn integration for model logging and loading
from sklearn.ensemble import RandomForestClassifier  # Ensemble learning model for classification tasks
from mlflow.models import infer_signature           # Auto-detects input/output schema from sample data

model = RandomForestClassifier(n_estimators=100, max_depth=10)  # Initialize 100 decision trees, each capped at depth 10
model.fit(X_train, y_train)             # Train the model on feature matrix X_train and target labels y_train

signature = infer_signature(X_train, model.predict(X_train))  # Infer column names, types, and shapes from training data and predictions

with mlflow.start_run():                # Create a new MLflow run context; all logging is grouped under this run's ID
    mlflow.sklearn.log_model(           # Log the trained model with environment, signature, and example to the artifact store
        sk_model=model,                 # The trained scikit-learn estimator to serialize and package
        artifact_path="fraud_detector", # Subdirectory name under the run's artifact root for model files
        signature=signature,            # Schema definition enforcing input validation and type coercion at inference time
        input_example=X_train.head(1)   # First row of training data saved as input_example.json for post-deployment testing
    )                                   # log_model completes; generates MLmodel, conda.yaml, requirements.txt, model.pkl

That log_model call writes the conda.yaml, requirements.txt, and the flavor metadata.

Step 2: Real-Time REST API (Zero Code)

Engineering doesn’t need to touch the training code. They grab the model from the registry and serve it with one command:

mlflow models serve -m "models:/fraud_detector/1" --port 5001 --host 0.0.0.0  # Serve model version 1 from the registry as a REST API on port 5001

Boom. REST API at localhost:5001/invocations. Send it a JSON payload:

curl -X POST http://localhost:5001/invocations \   # POST JSON input to the MLflow serving endpoint for real-time fraud scoring
  -H "Content-Type: application/json" \             # Tell the server the request body format is JSON
  -d '{"inputs": [[amount, time, merchant_score, ...]]}'  # Feature array matching the logged signature; each inner array is one transaction

The API returns {"predictions": [0]} — 0 means legit, 1 means fraud. No FastAPI code to write.

Step 3: Spark Batch Job (Also Zero Code)

The risk team needs to score 10 million rows every night in Spark. They load the exact same model as a Spark UDF:

import mlflow.pyfunc                   # MLflow's Python Function flavor for loading models in any Python environment
from pyspark.sql import SparkSession   # Spark SQL entry point for reading DataFrames and running queries

spark = SparkSession.builder.getOrCreate()              # Get or create the active SparkSession for DataFrame operations
transactions_df = spark.table("daily_transactions")     # Load the daily transactions table as a Spark DataFrame (millions of rows)

predict_udf = mlflow.pyfunc.spark_udf(                  # Register the MLflow model as a Spark UDF for distributed scoring
    spark,                                              # Active SparkSession used by the UDF for execution planning
    model_uri="models:/fraud_detector/1"                # URI pointing to model version 1 in the MLflow Model Registry
)                                                       # spark_udf returns a PySpark Column-compatible UDF

audited_df = transactions_df.withColumn("is_fraud", predict_udf())  # Apply the UDF to every row; result is 0 (legit) or 1 (fraud)
audited_df.write.saveAsTable("flagged_transactions")               # Persist the scored DataFrame as a new Hive/Spark table for compliance

Same model file. Same dependencies. Spark handles the parallelism.

The PEX Bonus Round

There’s one more trick: PEX (Python EXecutable). It bundles your model, all its dependencies, and the runtime into a single self-contained executable file.

import mlflow.pyfunc                   # MLflow's Python Function flavor for model loading and packaging

mlflow.pyfunc.build_pex(               # Build a self-contained PEX executable that bundles model + all dependencies
    model_uri="models:/fraud_detector/1",  # Model version 1 from the registry to embed inside the PEX file
    output_path="fraud_detector.pex"       # Output filename for the portable, single-file executable
)                                        # Result: fraud_detector.pex — no pip install needed on the target machine

Now you’ve got fraud_detector.pex — a single file that runs anywhere without installing anything:

./fraud_detector.pex                    # Execute the PEX file; it boots the model and starts serving predictions immediately

Ship it to an air-gapped server or a Docker scratch image. No dependency hell because everything’s baked in.

Why This Changes the Game

  • Data scientists use whatever they want. Operations teams never need to touch those libraries.
  • The conda.yaml and requirements.txt are generated during training. The Docker container and the Spark worker nodes both use the exact same versions.
  • The same models:/fraud_detector/1 URI works for a REST API, a Spark cluster, a PEX file, or a cron job.

No rewrites. No late-night debugging sessions comparing training and serving outputs.

The Bottom Line

MLflow Models doesn’t just solve the “works on my machine” problem. It solves the “works on my machine, the API server, the Spark cluster, and Bob’s laptop” problem.

Train once. Log with flavors. Deploy anywhere.

Comments

Leave a comment

What is 2 + 3?
Please solve the arithmetic question above to verify you're human.