Framework Flavors
MLflow integrates with 100+ ML and AI frameworks through its flavor system. Each flavor provides autologging, model saving/loading, and deployment support.
scikit-learn
mlflow.sklearn
Parameters via get_params(), metrics via score(), cloudpickle serialization.
PyTorch
mlflow.pytorch
Architecture, optimizer config, per-epoch metrics. Eager and TorchScript support.
TensorFlow
mlflow.tensorflow
Model summary, callback metrics, SavedModel format. TensorBoard integration.
XGBoost
mlflow.xgboost
Boosting params, feature importance, per-round eval metrics, early stopping.
LightGBM
mlflow.lightgbm
Native and sklearn API. Parameter logging and per-round metric tracking.
Spark MLlib
mlflow.spark
Pipeline stages, parameters, metrics. pyfunc wrapper for serving.
LangChain
mlflow.langchain
Chain tracing, prompt/completion logging, agent tool call capture.
OpenAI
mlflow.openai
API parameters, token usage, cost tracking, response logging.
Transformers
mlflow.transformers
Model config, training args, eval metrics. Automatic signature inference.
Common Integration Patterns
MLflow + Orchestrators (Airflow, Prefect, Dagster)
Use your orchestrator for scheduling; each step logs to MLflow. The orchestrator handles “when”; MLflow handles “what happened.”
MLflow + DVC (Data Versioning)
DVC versions datasets; MLflow tracks experiments. Log the DVC data hash as an MLflow parameter to link runs to dataset versions.
MLflow + Hyperparameter Tuning (Optuna, Ray Tune)
Each trial logs as a nested run. The parent captures the search space and best configuration.
Python
import optuna, mlflow
def objective(trial):
lr = trial.suggest_float("lr", 1e-5, 1e-1, log=True)
with mlflow.start_run(nested=True):
mlflow.log_param("lr", lr)
# Train and evaluate...
mlflow.log_metric("accuracy", accuracy)
return accuracy
with mlflow.start_run(run_name="optuna-search"):
study = optuna.create_study(direction="maximize")
study.optimize(objective, n_trials=100)Deployment Targets
Docker: Self-contained Kubernetes-ready images. AWS SageMaker: Direct deployment via mlflow.sagemaker.deploy(). Azure ML: Managed model serving. Databricks: Native integration with Unity Catalog governance.