Data Engineering

Reproducible ML Pipelines with DVC and MLflow

One of the most persistent challenges in Machine Learning engineering is the "it worked on my machine" problem. Unlike traditional software development, where inputs are fixed, ML models are trained on data. When that data changes, the model changes. Without strict versioning and tracking, teams often waste hours debugging why a model that performed well in testing fails in production, or why a specific improvement in accuracy cannot be reproduced. This article explores a robust workflow using Data Version Control (DVC) and MLflow to create fully reproducible training pipelines.

The Foundation: Data Versioning with DVC

Traditional version control systems like Git are optimized for text files, not large datasets. Data Version Control (DVC) extends Git's capabilities to handle large binary files, such as images, CSVs, and model artifacts. It does this by storing file hashes in Git while keeping the actual data in a remote storage location (like S3, GCS, or a local cache).

By integrating DVC into your workflow, you ensure that every model you train is associated with a specific, immutable snapshot of the training data. If you update your code but want to compare it against the exact same data version you used last week, you can simply check out that specific data commit.

# Initialize DVC in your repository
dvc init
dvc add data/raw_dataset.csv

# Track changes in the data pipeline
dvc add data/processed_dataset.csv -m
dvc run -n train_model -d data/processed_dataset.csv -o model.pkl python train.py

Experiment Tracking with MLflow

While DVC handles the data lineage, MLflow handles the experiment lineage. Training a model involves a complex matrix of hyperparameters, code versions, and environment variables. MLflow provides a unified interface to track these parameters, log metrics, and store artifacts.

In a standard training script, you wrap your logic in an MLflow run. This allows you to record the exact hyperparameters used (e.g., learning rate, number of layers) alongside the resulting performance metrics (e.g., accuracy, F1-score). Crucially, you can link this MLflow run to the specific DVC commit hash, creating a complete audit trail from data to result.

import mlflow
import dvc
import sys

def main():
    # Get the current DVC commit hash for reproducibility
    dvc_commit = dvc.repo.Repo().index.state.state_id

    with mlflow.start_run():
        # Log parameters
        mlflow.log_param("dvc_commit", dvc_commit)
        mlflow.log_param("learning_rate", 0.01)
        
        # Train model
        model = train_model(data_path="data/processed_dataset.csv")
        
        # Log metrics
        accuracy = evaluate(model)
        mlflow.log_metric("accuracy", accuracy)
        
        # Log artifacts
        mlflow.log_artifact("model.pkl")

if __name__ == "__main__":
    main()

Integrating into a CI/CD Pipeline

The true power of this stack is realized in Continuous Integration. When a developer pushes code, your CI system can automatically run the DVC pipeline to ensure the data is synced and then trigger MLflow experiments. If the new model's metrics degrade compared to the baseline, the pipeline can automatically fail the build. This shift-left approach to ML quality assurance catches issues early in the development cycle.

Conclusion

Building reproducible ML systems requires more than just writing clean code; it requires rigorous infrastructure for data and experiments. By combining DVC for deterministic data access and MLflow for granular experiment tracking, you eliminate the guesswork from model deployment. This stack ensures that any result you claim is backed by verifiable data and code, fostering trust and reliability in your machine learning operations.

Share: