DEV Community

Ivan Zykov
Ivan Zykov

Posted on

I Squeezed an Entire MLOps Pipeline into 10 Lines of YAML

#mlops #gitlab #cicd #machinelearning

Every ML project I worked on in GitLab had the same problem: a bloated .gitlab-ci.yml with hand-rolled MLflow integration, custom validation scripts, and manual model registration. Copy it to the next project, tweak the paths, fix the bugs you already fixed last time. By the fifth project you don't remember which config has the working version of MLFLOW_RUN_ID passthrough between jobs.

So I built a GitLab CI/CD component that replaces all of that with 10 lines of YAML.

Before vs After

Here's what a typical MLOps pipeline looked like before — and this is the shortened version:

stages: [validate, train, evaluate, register]

validate-data:
  stage: validate
  image: python:3.12
  script:
    - pip install pandas great_expectations
    - python scripts/validate.py --data data/train.csv --check-nulls --threshold 0.05
  artifacts:
    paths: [validation_report.json]

train-model:
  stage: train
  image: python:3.12
  variables:
    MLFLOW_TRACKING_URI: "${CI_API_V4_URL}/projects/${CI_PROJECT_ID}/ml/mlflow"
  script:
    - pip install mlflow scikit-learn pandas
    - python scripts/train.py --data data/train.csv
    - echo "MLFLOW_RUN_ID=$(cat run_id.txt)" >> train.env
  artifacts:
    reports:
      dotenv: train.env
    paths: [model/, metrics.json]

evaluate-model:
  stage: evaluate
  image: python:3.12
  needs: [{job: train-model, artifacts: true}]
  variables:
    MLFLOW_TRACKING_URI: "${CI_API_V4_URL}/projects/${CI_PROJECT_ID}/ml/mlflow"
  script:
    - pip install mlflow
    - python scripts/evaluate.py --run-id $MLFLOW_RUN_ID --threshold 0.85
    - echo "EVAL_PASSED=$(cat eval_result.txt)" >> evaluate.env
  artifacts:
    reports:
      dotenv: evaluate.env

register-model:
  stage: register
  image: python:3.12
  needs: [{job: train-model, artifacts: true}, {job: evaluate-model, artifacts: true}]
  rules:
    - if: $EVAL_PASSED == "true"
  variables:
    MLFLOW_TRACKING_URI: "${CI_API_V4_URL}/projects/${CI_PROJECT_ID}/ml/mlflow"
  script:
    - pip install mlflow
    - python scripts/register.py --run-id $MLFLOW_RUN_ID --model-name my-model
Enter fullscreen mode Exit fullscreen mode

And this doesn't even include DVC, pip caching, retry logic, or error handling. Each project also had its own validate.py, evaluate.py, register.py — each with its own implementation of auto_configure_mlflow, its own argument parsing, its own bugs.

Now the same thing:

stages: [validate, train, evaluate, register]

include:
  - component: gitlab.com/netOpyr/gitlab-mlops-component/full-pipeline@1.0.0
    inputs:
      model_name: wine-classifier
      training_script: scripts/train.py
      training_args: '--data data/train.csv --test-data data/test.csv'
      data_path: data/train.csv
      framework: sklearn
      metric_name: accuracy
      min_threshold: '0.85'
Enter fullscreen mode Exit fullscreen mode

These 10 lines give you 4 jobs:

validate --> train --> evaluate --> register
   |           |          |            |
 schema      MLflow    accuracy    Model Registry
 nulls       autolog   >= 0.85    (if eval passed)
 drift       metrics   vs prod
Enter fullscreen mode Exit fullscreen mode

All the boilerplate scripts now live inside the component. You only write the training script.

What Each Stage Does

validate checks your data before training starts: schema validation (are all columns present?), null ratio per column (default threshold 5%), and optionally data drift detection via Evidently. Supports Great Expectations suites and custom Python check scripts too.

train wraps your training script in an MLflow session. It auto-configures the tracking URI from GitLab CI variables, creates an experiment and run, enables autolog for your framework (sklearn, PyTorch, TensorFlow, XGBoost, LightGBM), and passes MLFLOW_RUN_ID to your script via environment variable. Your script stays a normal Python file — it works locally and in Jupyter just the same.

evaluate pulls metrics from MLflow and runs them through quality gates. Gate 1: absolute threshold (e.g. accuracy >= 0.85). Gate 2 (optional): comparison with the current production model from Model Registry. Supports higher_is_better: false for loss metrics.

register pushes the model to GitLab Model Registry with metadata: alias (staging by default), commit SHA, pipeline ID, metrics. Works on all GitLab tiers — on Free, alias assignment silently falls back to tags.

DVC Integration

If your data lives in S3/MinIO:

include:
  - component: .../train@1.0.0
    inputs:
      training_script: scripts/train.py
      model_name: my-model
      dvc_enabled: true
      dvc_remote: minio
      dvc_files: 'data/train.csv.dvc data/test.csv.dvc'
      dvc_push: true
      dvc_push_paths: 'model/'
Enter fullscreen mode Exit fullscreen mode

The component installs DVC, pulls data before training, and pushes artifacts back after. Credentials go through CI/CD variables: AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_ENDPOINT_URL.

When 10 Lines Aren't Enough

For more complex setups you can include each stage separately:

include:
  - component: .../validate@1.0.0
    inputs:
      data_path: data/train.csv
      enable_drift: true
      reference_data_path: data/reference.csv

  - component: .../train@1.0.0
    inputs:
      training_script: scripts/train.py
      model_name: my-model
      image_suffix: pytorch-gpu
      framework: pytorch
      tags: ["gpu"]

  - component: .../evaluate@1.0.0
    inputs:
      model_name: my-model
      metric_name: val_loss
      min_threshold: '0.1'
      higher_is_better: false

  - component: .../register@1.0.0
    inputs:
      model_name: my-model
      alias: staging
Enter fullscreen mode Exit fullscreen mode

This way you get per-stage GPU runners, custom images, and conditional execution. You can also train multiple models in parallel using the as parameter to give unique names to jobs.

Framework Support

Each framework has a dedicated Docker image selected via image_suffix:

Suffix Frameworks GPU
sklearn scikit-learn, matplotlib No
boosting XGBoost, LightGBM, scikit-learn No
pytorch PyTorch (CPU) No
pytorch-gpu PyTorch + CUDA 12.4 Yes
tensorflow TensorFlow (CPU) No
tensorflow-gpu TensorFlow + CUDA 12.4 Yes

All images include Python 3.12, MLflow, and pandas. Need extra dependencies? Set requirements_file: requirements.txt or bring your own image via image_registry_base.

Try It

Two options:

  1. Fork the example project — a wine classifier with 3 files total. Just create an access token with API scope and add it as MLOPS_ACCESS_TOKEN in CI/CD variables.
  2. Add the component to your existing project. Drop your training script into scripts/ and configure the inputs.

The component is published in the GitLab CI/CD Catalog.

What's Next

Coming soon: BuildKit-based image builds, retry logic for flaky MLflow requests, and GitLab Environments integration.

Found a bug or missing a feature? Open an issue or MR.

Links:

Top comments (0)