How do I set up a training workflow platform for reproducible experiments with dataset versioning and GPU support?

A solid setup usually has 5 pieces:

1. Code versioning
2. Dataset/version tracking
3. Environment reproducibility
4. Experiment tracking
5. GPU-capable execution/orchestration

Here’s a practical blueprint you can use.

---

1) Pick a platform stack

A common, reliable stack looks like this:

• Git for code
• DVC or lakeFS for dataset versioning
• Docker or Conda + lockfiles for environment reproducibility
• MLflow or Weights & Biases for experiment tracking
• Kubernetes + NVIDIA GPU nodes or a single GPU machine with Docker for execution
• Optional workflow engine: Airflow, Prefect, Kubeflow, or Metaflow

If you want the simplest strong setup:

• Git + DVC + Docker + MLflow
• run jobs on a GPU server or Kubernetes GPU cluster

---

2) Organize the repo for reproducibility

A clean project structure helps a lot:

project/
├── data/
│   ├── raw/
│   ├── processed/
│   └── .gitignore
├── models/
├── notebooks/
├── src/
│   ├── train.py
│   ├── evaluate.py
│   ├── preprocess.py
│   └── config.py
├── configs/
│   ├── train.yaml
│   └── data.yaml
├── dvc.yaml
├── params.yaml
├── Dockerfile
├── requirements.txt  # or environment.yml / poetry.lock
└── README.md

Keep:

• code in src/
• configs in configs/
• raw data out of Git
• artifacts in tracked storage via DVC/MLflow

---

3) Add dataset versioning with DVC

Install and initialize

pip install dvc
git init
dvc init

Track a dataset

dvc add data/raw/train.csv
git add data/raw/train.csv.dvc .gitignore
git commit -m "Track training data with DVC"

Connect remote storage

Use S3, GCS, Azure Blob, SSH, or local network storage.

Example with S3:

dvc remote add -d myremote s3://my-bucket/dvcstore
dvc remote modify myremote region us-east-1
dvc push

Now the repo tracks the pointer to the dataset version, while the actual data stays in object storage.

Why this matters

Anyone can reproduce the same training run by checking out:

• the same Git commit
• the same DVC data version
• the same code/environment

---

4) Make the environment reproducible

Best option: Docker

Create a Docker image with exact dependencies.

Example Dockerfile:

FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04

WORKDIR /app

RUN apt-get update && apt-get install -y python3 python3-pip

COPY requirements.txt .
RUN pip3 install --no-cache-dir -r requirements.txt

COPY . .
CMD ["python3", "src/train.py"]

If you need GPU support, make sure:

• host has NVIDIA drivers
• install NVIDIA Container Toolkit
• run container with GPU access

Example:

docker run --gpus all -it my-training-image

Alternative: Conda/Poetry

Good for local dev, but Docker is better for strict reproducibility.

---

5) Track experiments and metrics

Use MLflow or Weights & Biases.

Example with MLflow

pip install mlflow

In your training script:

import mlflow

with mlflow.start_run():
    mlflow.log_param("lr", 0.001)
    mlflow.log_param("batch_size", 32)
    mlflow.log_metric("accuracy", 0.91)
    mlflow.log_artifact("models/model.pt")

You can then compare:

• hyperparameters
• metrics
• artifacts
• run metadata

This is essential for reproducibility and auditability.

---

6) Define the pipeline

Use DVC pipelines or a workflow engine.

Example dvc.yaml:

stages:
  preprocess:
    cmd: python src/preprocess.py data/raw/train.csv data/processed/train.parquet
    deps:
      - src/preprocess.py
      - data/raw/train.csv
    outs:
      - data/processed/train.parquet

  train:
    cmd: python src/train.py data/processed/train.parquet models/model.pkl
    deps:
      - src/train.py
      - data/processed/train.parquet
    outs:
      - models/model.pkl
    params:
      - train.lr
      - train.batch_size

Run:

dvc repro

This gives you a reproducible, dependency-aware pipeline.

---

7) Enable GPU training

On a single machine

1. Install NVIDIA driver
2. Install CUDA-compatible PyTorch/TensorFlow
3. Use Docker with --gpus all

Check GPU availability in Python:

import torch
print(torch.cuda.is_available())
print(torch.cuda.get_device_name(0))

On Kubernetes

If you use K8s:

• install NVIDIA device plugin
• request GPU resources in your job spec

Example:

resources:
  limits:
    nvidia.com/gpu: 1

This allows scheduled GPU workloads across a cluster.

---

8) Version everything needed for a run

For each experiment, store:

• Git commit hash
• DVC dataset version
• Docker image tag/digest
• hyperparameters
• random seed
• code entrypoint
• metrics and model artifact

Example run metadata:

{
  "git_commit": "a1b2c3d",
  "data_version": "dvc:9f8e7d",
  "image": "training:1.4.2",
  "seed": 42,
  "lr": 0.001,
  "batch_size": 32
}

This is the minimum needed to reproduce a result later.

---

9) Automate with CI/CD

Use GitHub Actions, GitLab CI, or Jenkins to:

• lint code
• run unit tests
• validate configs
• build Docker image
• optionally run small CPU smoke tests
• push artifacts

Example CI checks:

• pytest
• ruff / flake8
• dvc doctor
• docker build

You can also trigger training on merge or on tagged releases.

---

10) Recommended minimal setup

If you want something realistic and not too heavy:

Local/small team

• GitHub
• DVC + S3
• Docker
• MLflow
• one GPU server
• GitHub Actions for CI

Larger team / production

• Git
• DVC or lakeFS
• MLflow/W&B
• Kubernetes with GPU nodes
• Prefect/Kubeflow/Argo for orchestration
• centralized artifact store (S3/GCS/MinIO)

---

11) A simple end-to-end workflow

1. Commit code to Git
2. Update data via DVC
3. Build Docker image
4. Launch training on GPU
5. Log metrics/artifacts to MLflow
6. Save model + run metadata
7. Reproduce by checking out same commit + DVC version + Docker image

---

12) Best practices

• Set random seeds
• Use deterministic ops where possible
• Log library versions (pip freeze, conda env export)
• Keep configs in YAML, not hardcoded
• Store raw/processed data separately
• Version both preprocessing and training code
• Use immutable Docker image tags or digests
• Test on a small dataset before large GPU runs

---

Suggested starter stack

If you want a strong default choice:

• Git + DVC + Docker + MLflow
• storage: S3 or MinIO
• compute: one NVIDIA GPU machine or Kubernetes GPU cluster

That will give you reproducible experiments, dataset versioning, and GPU support without too much complexity.

If you want, I can also provide:

1. a sample repo template,
2. a Dockerfile + DVC + MLflow example, or
3. a Kubernetes-based setup for multi-GPU training.