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_functionflavor. - 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.yamlandrequirements.txtare generated during training. The Docker container and the Spark worker nodes both use the exact same versions. - The same
models:/fraud_detector/1URI 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