Hardware and framework version tracking in MLOps - Commands & Configuration

Introduction

When you train machine learning models, it is important to know exactly what hardware and software versions were used. This helps you reproduce results and fix problems later. Tracking hardware and framework versions automatically saves time and avoids confusion.

When you want to record the GPU model and driver version used for training a model.

When you need to log the exact version of TensorFlow or PyTorch your code ran on.

When you want to compare model performance across different hardware setups.

When you need to share your model training environment details with teammates.

When you want to ensure your model can be reproduced months later with the same setup.

Commands

Install MLflow, a tool that helps track machine learning experiments including hardware and framework versions.

Terminal

pip install mlflow

Expected Output

Collecting mlflow

  Downloading mlflow-2.4.1-py3-none-any.whl (18.7 MB)

Installing collected packages: mlflow

Successfully installed mlflow-2.4.1

Run a Python script that logs hardware and framework versions to MLflow for tracking.

Terminal

python track_versions.py

Expected Output

2024/06/01 12:00:00 INFO mlflow.tracking.fluent: Experiment with name 'HardwareFrameworkTracking' does not exist. Creating a new experiment.

2024/06/01 12:00:00 INFO mlflow.tracking.fluent: Experiment created with ID '1'

2024/06/01 12:00:00 INFO mlflow.tracking.fluent: Run started with run_id '1234567890abcdef'

Logged parameter: gpu_name = NVIDIA GeForce RTX 3080

Logged parameter: cuda_version = 11.8

Logged parameter: pytorch_version = 2.0.1

Logged parameter: python_version = 3.11.4

Run ended with status FINISHED

Key Concept

If you remember nothing else from this pattern, remember: automatically logging hardware and framework versions ensures your ML experiments are reproducible and understandable later.

Code Example

MLOps

import mlflow
import platform
import torch
import subprocess

# Function to get GPU name

def get_gpu_name():
    try:
        result = subprocess.run(['nvidia-smi', '--query-gpu=name', '--format=csv,noheader'],
                                stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=True)
        return result.stdout.strip()
    except Exception:
        return 'No GPU detected'

# Function to get CUDA version

def get_cuda_version():
    try:
        result = subprocess.run(['nvcc', '--version'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=True)
        for line in result.stdout.splitlines():
            if 'release' in line:
                return line.split('release')[-1].strip().split(',')[0]
        return 'Unknown CUDA version'
    except Exception:
        return 'CUDA not installed'

mlflow.set_experiment('HardwareFrameworkTracking')
with mlflow.start_run():
    gpu_name = get_gpu_name()
    cuda_version = get_cuda_version()
    pytorch_version = torch.__version__
    python_version = platform.python_version()

    mlflow.log_param('gpu_name', gpu_name)
    mlflow.log_param('cuda_version', cuda_version)
    mlflow.log_param('pytorch_version', pytorch_version)
    mlflow.log_param('python_version', python_version)

    print('Logged hardware and framework versions to MLflow')

Common Mistakes

Not logging hardware or framework versions at all.

Without this info, you cannot reproduce or debug your model training environment.

Always include code to log hardware and framework versions as parameters in your experiment tracking.

Manually writing hardware info in notes instead of automating it.

Manual notes can be forgotten or inaccurate, causing confusion.

Use code to detect and log hardware and software versions automatically.

Summary

Install MLflow to track machine learning experiments.

Write a Python script to automatically detect and log hardware and framework versions.

Run the script to save this info in MLflow for reproducibility and debugging.