Robust AI Pipelines with Data Versioning Using DVC and MLflow

Opening: Why does the same dataset feel different each time you run a notebook?

It happens more often than you think. You open a Jupyter notebook, pull the latest code from master, and the model you trained yesterday suddenly gives a worse accuracy. The culprit is usually a hidden change in the raw files – a new row added, a column renamed, or a preprocessing script tweaked. When the data source moves without a clear version marker, you lose the ability to explain the shift, roll back, or let a teammate reproduce your results. This post walks through a practical way to lock down data and model artefacts using DVC or MLflow, add a model registry, watch for drift, and tie everything together in a CI/CD pipeline.

Background: Why data versioning matters now more than ever

Compliance frameworks increasingly demand traceability of every artefact that influences a decision‑making model. Finance, healthcare, and regulated SaaS products must show which data version produced a given prediction, and they must be able to revert if a problem is discovered. At the same time, teams are moving toward micro‑service architectures where models are deployed as containers and updated several times a day. Without a reliable way to version datasets, you end up with a tangled web of ad‑hoc scripts and manual copy‑pastes. The cost of a broken pipeline is not just a missed deadline – it can be a legal exposure.

Both DVC and MLflow were created to address this gap. DVC treats data files like code, storing hashes in Git and pushing the actual blobs to a remote storage (S3, Azure, GCS, etc.). MLflow adds a model registry on top, allowing you to promote models from “staging” to “production” with a single API call. When you combine them, you get a reproducible end‑to‑end workflow that can survive team turnover and infrastructure changes.

Technical deep dive: Architecture and a minimal example

The core idea is to keep three things in sync:

• Data files (raw, processed, features) – versioned with DVC.
• Model artefacts and metadata – logged to MLflow.
• Deployment descriptors – stored in Git and triggered by a CI pipeline.

Below is a stripped‑down example that shows how to train a scikit‑learn model, log the data version, and register the model. The script assumes you have a DVC remote configured (for example an S3 bucket) and an MLflow tracking server reachable at http://mlflow.mycompany.com.

# train.py
import os
import dvc.api
import mlflow
import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score

# 1. Pull a specific version of the dataset using DVC API
# The URL points to the data file in the DVC repo, the rev can be a tag, branch or commit hash
DATA_URL = "https://github.com/yourorg/ai-data.git/data/transactions.csv"
DATA_REV = os.getenv("DATA_REV", "main")
with dvc.api.open(
    path=DATA_URL,
    rev=DATA_REV,
    mode="r",
    repo="https://github.com/yourorg/ai-data.git",
) as fd:
    df = pd.read_csv(fd)

X = df.drop(columns=["target"])
y = df["target"]
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

mlflow.set_tracking_uri("http://mlflow.mycompany.com")
mlflow.set_experiment("transaction‑fraud")

with mlflow.start_run() as run:
    # Log the data version as a tag – this ties the model to the exact data snapshot
    mlflow.set_tag("data_rev", DATA_REV)
    model = RandomForestClassifier(n_estimators=100, random_state=42)
    model.fit(X_train, y_train)
    preds = model.predict(X_test)
    acc = accuracy_score(y_test, preds)
    mlflow.log_metric("accuracy", acc)
    # Log the model itself
    mlflow.sklearn.log_model(model, "model")
    # Register the model (creates a new version in the registry)
    mlflow.register_model(f"runs:/{run.info.run_id}/model", "FraudModel")
    print(f"Model logged with accuracy {acc:.4f}")

Key points in the snippet:

• The dvc.api.open call fetches the data at the exact commit identified by DATA_REV. If the remote data changes, the hash will differ and the script will automatically pull the new version.
• MLflow tags are used to store the data revision. Later you can query the registry for all models built on a specific data snapshot.
• Registering the model creates a versioned entry in the MLflow Model Registry. You can promote that version to “Production” with a UI click or an API call.

Trade‑offs:

• Storage cost: DVC stores each data version as a separate object in the remote. Large raw files can quickly fill up an S3 bucket. Strategies like data deduplication, compression, or storing only delta files help keep costs manageable.
• Metadata complexity: You now have two sources of truth – Git for code, DVC for data, and MLflow for models. Keeping them consistent requires disciplined branching and CI checks.
• Performance overhead: Pulling a new data version adds network latency, especially for multi‑GB files. Caching the latest version locally and only pulling when the hash changes mitigates the impact.

Common pitfalls: What tends to go wrong

Even with a solid blueprint, teams hit snags early on. Here are the most frequent issues and how to avoid them:

• Forgetting to lock the data revision: If you run mlflow.start_run without setting the data_rev tag, you lose the link between model and dataset. Make the tag addition part of a shared utility function.
• Large monolithic data files: Storing a 50 GB CSV in DVC leads to slow checkout and high egress costs. Break the data into logical partitions (by date or region) and version each partition separately.
• Remote storage permission drift: When a new team joins, they often lack read access to the DVC remote. Automate IAM policies as part of your onboarding scripts.
• Model registry naming collisions: Registering many models under the same name without a clear naming convention creates confusion. Prefix models with the domain (e.g., FraudModel, ChurnModel) and include the data tag in the description.
• CI pipelines failing silently on data pull errors: If DVC cannot fetch the data, the job may still succeed but produce an empty model. Add explicit checks for file existence and non‑zero size after dvc pull.

Practical implementation guide: Step‑by‑step for a reproducible pipeline

Below is a concise checklist you can copy into a README or a team wiki. Adjust paths and tool versions to your environment.

1. Initialize the repository
   • Run git init and git remote add origin ….
   • Install DVC (>=2.30) and configure a remote: dvc remote add -d storage s3://my‑bucket/dvc then dvc remote modify storage access_key_id $AWS_ACCESS_KEY_ID and dvc remote modify storage secret_access_key $AWS_SECRET_ACCESS_KEY.
2. Version your raw data
   • Place the raw CSV in data/raw/transactions.csv.
   • Run dvc add data/raw/transactions.csv – this creates a .gitignore entry and a .dvc file with the hash.
   • Commit both .dvc and .gitignore: git add . && git commit -m "Add raw transaction data".
3. Set up MLflow tracking
   • Deploy an MLflow server (Docker image mlflow:latest) behind an internal load balancer.
   • Create an experiment via UI or mlflow experiments create -n transaction‑fraud.
4. Write the training script
   • Use the code snippet from the technical deep dive as a template.
   • Parameterise the data revision with an environment variable (e.g., DATA_REV).
5. Configure CI/CD
   • In your GitHub Actions workflow, add steps: checkout, setup python, pip install -r requirements.txt dvc[aws] mlflow, dvc pull (this fetches the exact data version), and finally python train.py.
   • After a successful run, add a step that queries the MLflow registry for the latest model version and builds a Docker image with mlflow models serve as the entry point.
   • Deploy the image to your Kubernetes cluster using a Helm chart that references the model version tag.
6. Monitor drift
   • Schedule a nightly job that loads the latest production data, computes feature statistics, and compares them to the statistics stored in the model’s tags (you can log a JSON summary with mlflow.log_dict).
   • If a statistical test exceeds a threshold, raise an alert and create a Git issue automatically.

Following this checklist gives you a pipeline that can be triggered by a pull request, audited by a compliance auditor, and rolled back with a single git revert + dvc checkout.

Closing thoughts: Keep the system simple and observable

The biggest temptation is to add every possible feature – automatic hyper‑parameter sweeps, model explainability dashboards, and multi‑cloud data backends – before the basic workflow is stable. In practice, a lean pipeline that reliably records data hash, model version, and a few key metrics is far more valuable. Once the core loop works, you can extend it incrementally, always asking “does this addition keep the reproducibility guarantee intact?”. Remember to document the naming conventions for data tags and model versions; a well‑written README often saves more time than any automation.

Ready to try it out?

Grab the starter repo at akkistech.com, spin up the MLflow server, and follow the checklist. The first version may feel a bit heavy, but the confidence you gain when you can roll back a model to the exact data snapshot is priceless.