ā† Back to DevBytes
Docker July 07, 2026 13 min read

Docker with GPU ML Workloads: Production Guide

What Are Docker GPU ML Workloads?

Docker GPU ML workloads refer to the practice of containerizing machine learning applications that require NVIDIA GPU acceleration, then deploying them consistently across development, staging, and production environments. At its core, this involves installing the NVIDIA Container Toolkit on the host machine, selecting a base image with CUDA libraries baked in, and granting the container access to the host GPUs via the --gpus runtime flag. The result is a portable, reproducible environment where your PyTorch, TensorFlow, or JAX model trains or serves inferences identically on a single workstation, a Kubernetes cluster, or a cloud GPU instance.

Why GPU Containerization Matters in Production

šŸš€ Deploy your AI agent in 10 minutes

Managed Hermes hosting. Zero DevOps. 100M tokens/mo included.

Try it free →

Without containers, GPU ML workflows suffer from dependency hell: mismatched CUDA versions, conflicting cuDNN libraries, driver incompatibilities, and Python package drift. Containers solve this by packaging the entire runtime stackā€”OS libraries, CUDA toolkit, cuDNN, Python, and your model dependenciesā€”into a single immutable artifact. This means:

Prerequisites and Host Setup

Before you can run GPU containers, the host machine needs three things: an NVIDIA GPU, a compatible driver, and the NVIDIA Container Toolkit. Here's how to install the toolkit on an Ubuntu 22.04 host:

# 1. Add NVIDIA package repositories
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg

curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
  sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#' | \
  sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list

# 2. Install the toolkit
sudo apt-get update
sudo apt-get install -y nvidia-container-toolkit

# 3. Configure Docker to use the NVIDIA runtime
sudo nvidia-ctk runtime configure --runtime=docker
sudo systemctl restart docker

# 4. Verify GPU access from inside a container
docker run --rm --gpus all nvidia/cuda:12.4.1-base-ubuntu22.04 nvidia-smi

The final verification command should print your GPU inventory inside the containerā€”proof that the CUDA driver and device files are correctly mounted. If you see nvidia-smi output, you're ready to build GPU containers.

Building Your First GPU-Ready Docker Image

The foundation of any GPU ML image is a CUDA base image. NVIDIA publishes several variants under the nvidia/cuda namespace. For production, prefer the -runtime or -devel flavors depending on whether you need compilation tools. Here's a production-oriented Dockerfile for a PyTorch model server:

# ---- Build Stage ----
FROM nvidia/cuda:12.4.1-devel-ubuntu22.04 AS builder

ENV DEBIAN_FRONTEND=noninteractive \
    PYTHONUNBUFFERED=1

RUN apt-get update && apt-get install -y --no-install-recommends \
    python3.11 \
    python3.11-venv \
    python3-pip \
    git \
    curl \
    && rm -rf /var/lib/apt/lists/*

# Create virtual environment
RUN python3.11 -m venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH"

# Install PyTorch with CUDA support and lock versions
COPY requirements.txt /tmp/requirements.txt
RUN pip install --no-cache-dir -r /tmp/requirements.txt

# ---- Runtime Stage ----
FROM nvidia/cuda:12.4.1-runtime-ubuntu22.04

ENV DEBIAN_FRONTEND=noninteractive \
    PYTHONUNBUFFERED=1

# Install only runtime dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
    python3.11 \
    python3.11-venv \
    libgomp1 \
    && rm -rf /var/lib/apt/lists/*

# Copy virtual environment from builder
COPY --from=builder /opt/venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH"

# Copy application code
WORKDIR /app
COPY ./src /app/src
COPY ./config /app/config

# Create non-root user for security
RUN groupadd -r appgroup && useradd -r -g appgroup -d /app -s /sbin/nologin appuser
RUN chown -R appuser:appgroup /app /opt/venv
USER appuser

EXPOSE 8080

HEALTHCHECK --interval=30s --timeout=10s --retries=3 \
  CMD curl -f http://localhost:8080/health || exit 1

ENTRYPOINT ["python3.11", "-m", "src.server"]

The corresponding requirements.txt might look like:

torch==2.4.0
torchvision==0.19.0
transformers==4.42.4
fastapi==0.112.0
uvicorn[standard]==0.30.6
pydantic==2.8.2
nvidia-ml-py==12.560.0
prometheus-client==0.20.0

This multi-stage build pattern is critical for production: the heavyweight CUDA development libraries, git, and build tools stay in the builder stage and never appear in the final runtime image. The runtime stage copies only the compiled virtual environment and your application code, shrinking the image by hundreds of megabytes and reducing the attack surface.

Running GPU Containers with Docker Compose

For local development or single-node production, Docker Compose simplifies GPU container orchestration. Here's a complete docker-compose.yml that runs a model inference service alongside a Redis cache and a Prometheus metrics collector:

version: "3.8"

services:
  inference:
    build:
      context: .
      dockerfile: Dockerfile
      args:
        CUDA_VERSION: "12.4.1"
    image: myregistry/inference-service:${IMAGE_TAG:-latest}
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [gpu]
    environment:
      - CUDA_VISIBLE_DEVICES=0
      - MODEL_PATH=/models/checkpoint.pt
      - BATCH_SIZE=32
      - MAX_CONCURRENT_REQUESTS=100
    volumes:
      - /mnt/models:/models:ro
      - inference-logs:/var/log/inference
    ports:
      - "8080:8080"
    healthcheck:
      test: ["CMD", "curl", "-f", "http://localhost:8080/health"]
      interval: 15s
      timeout: 5s
      retries: 3
    restart: unless-stopped
    logging:
      driver: json-file
      options:
        max-size: "10m"
        max-file: "3"

  redis:
    image: redis:7.4-alpine
    command: redis-server --maxmemory 2gb --maxmemory-policy allkeys-lru
    volumes:
      - redis-data:/data
    restart: unless-stopped

  prometheus-exporter:
    image: nvidia/dcgm-exporter:3.3.5
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]
    command: ["-f", "/etc/dcgm-exporter/default-counters.csv"]
    ports:
      - "9400:9400"
    restart: unless-stopped

volumes:
  inference-logs:
  redis-data:

Key production details in this compose file:

Handling GPU Memory and Resource Limits

GPU memory is a precious resource. A container can easily exhaust VRAM and cause out-of-memory errors that cascade across co-located services. In production, you must enforce limits:

# Example: run with explicit GPU memory and compute constraints
docker run --rm \
  --gpus '"device=0,capabilities=utility,compute"' \
  --memory=8g \
  --memory-swap=8g \
  --shm-size=1g \
  --ulimit memlock=-1 \
  myregistry/inference-service:latest

The --shm-size=1g flag is especially important for PyTorch data loaders that use shared memory for inter-process communication. Without it, multi-worker dataloaders crash with "insufficient shared memory" errors. The --ulimit memlock=-1 removes the lockable memory ceiling, which CUDA requires for GPU memory allocation. Inside your application code, proactively manage VRAM:

import torch
import os

def configure_gpu_memory():
    """Apply production-grade GPU memory settings."""
    device_count = torch.cuda.device_count()
    visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", "all")
    
    for i in range(device_count):
        # Set a hard fraction of total VRAM for this process
        torch.cuda.set_per_process_memory_fraction(0.85, device=i)
        
        # Enable aggressive caching allocator cleanup
        torch.cuda.memory.set_per_process_memory_fraction(0.85, device=i)
    
    # Use memory-efficient attention if available
    if hasattr(torch.nn.functional, "scaled_dot_product_attention"):
        torch.backends.cuda.enable_mem_efficient_sdp(True)
    
    # Disable TF32 if precision matters more than throughput
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True

    print(f"GPU memory configured for {device_count} device(s)")
    return device_count

Model Loading and Caching Strategies

Production GPU containers often start and stop frequently during rolling deployments. Downloading a multi-gigabyte model on every cold start is unacceptable. Instead, bake small models into the image and mount large ones from persistent storage:

# Inside your inference server startup code
import torch
import os
from pathlib import Path

class ModelManager:
    def __init__(self, model_dir: str = "/models"):
        self.model_dir = Path(model_dir)
        self.cache_dir = Path("/app/model_cache")
        self.cache_dir.mkdir(exist_ok=True)
    
    def load_model(self, model_id: str):
        """Load model with tiered caching: local cache ā†’ mounted volume ā†’ HuggingFace."""
        cache_path = self.cache_dir / model_id
        volume_path = self.model_dir / model_id
        
        # Tier 1: Check local warm cache (fastest)
        if cache_path.exists():
            return torch.load(cache_path / "checkpoint.pt", map_location="cuda")
        
        # Tier 2: Check mounted volume (fast, updated externally)
        if volume_path.exists():
            checkpoint = torch.load(volume_path / "checkpoint.pt", map_location="cuda")
            # Populate cache for next restart
            import shutil
            shutil.copytree(volume_path, cache_path, dirs_exist_ok=True)
            return checkpoint
        
        # Tier 3: Pull from remote registry (slow, only on cache miss)
        from huggingface_hub import snapshot_download
        snapshot_download(repo_id=model_id, local_dir=str(cache_path))
        return torch.load(cache_path / "checkpoint.pt", map_location="cuda")

For frequently updated models (like recommendation systems that retrain hourly), store the checkpoint on an object store (S3, GCS) and use an init container pattern in Kubernetes, or a pre-start script in Docker Compose, to pull the latest version before the main container starts.

Production-Grade Inference Server Example

Below is a minimal but production-hardened FastAPI inference server that handles batching, graceful shutdown, and health checks correctly. Place this in src/server.py:

import torch
import torch.nn.functional as F
from fastapi import FastAPI, Request
from pydantic import BaseModel, Field
from contextlib import asynccontextmanager
import signal
import sys
import time
import logging
from typing import List, Optional

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("inference-server")

# ---- Configuration ----
class Config:
    MODEL_PATH: str = "/models/checkpoint.pt"
    BATCH_SIZE: int = 32
    MAX_WAIT_SECONDS: float = 0.05  # dynamic batching window
    MAX_CONCURRENT: int = 100
    
config = Config()

# ---- Request/Response Schemas ----
class PredictRequest(BaseModel):
    input_ids: List[int] = Field(..., min_length=1, max_length=2048)
    temperature: Optional[float] = Field(0.7, ge=0.1, le=2.0)
    
class PredictResponse(BaseModel):
    logits: List[float]
    inference_time_ms: float
    model_version: str

# ---- Lifecycle Management ----
model: Optional[torch.nn.Module] = None
model_version: str = "unknown"
request_semaphore = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    global model, model_version, request_semaphore
    import asyncio
    
    logger.info("Loading model from %s", config.MODEL_PATH)
    model = torch.load(config.MODEL_PATH, map_location="cuda")
    model.eval()
    model_version = getattr(model, "version", "v1.0.0")
    request_semaphore = asyncio.Semaphore(config.MAX_CONCURRENT)
    
    # Warm up GPU with a dummy forward pass
    with torch.no_grad():
        dummy_input = torch.randint(0, 1000, (1, 128)).cuda()
        _ = model(dummy_input)
        torch.cuda.synchronize()
    
    logger.info("Model loaded. Version: %s. Ready.", model_version)
    yield
    
    logger.info("Shutting down. Draining requests...")
    # Let in-flight requests finish
    await asyncio.sleep(5)
    del model
    torch.cuda.empty_cache()
    logger.info("Shutdown complete.")

app = FastAPI(lifespan=lifespan, title="GPU Inference Service", version="2.0.0")

# ---- Middleware ----
@app.middleware("http")
async def add_metrics_header(request: Request, call_next):
    start = time.perf_counter()
    response = await call_next(request)
    elapsed_ms = (time.perf_counter() - start) * 1000
    response.headers["X-Inference-Time-Ms"] = f"{elapsed_ms:.2f}"
    return response

# ---- Endpoints ----
@app.get("/health")
async def health_check():
    if model is None:
        return {"status": "not_ready", "gpu": torch.cuda.is_available()}
    # Verify GPU memory is healthy
    free_memory = torch.cuda.mem_get_info()[0] / 1e9
    return {
        "status": "healthy",
        "model_version": model_version,
        "gpu_free_memory_gb": round(free_memory, 2),
        "cuda_devices": torch.cuda.device_count()
    }

@app.post("/predict", response_model=PredictResponse)
async def predict(request: PredictRequest):
    async with request_semaphore:
        start = time.perf_counter()
        with torch.no_grad():
            input_tensor = torch.tensor([request.input_ids]).cuda()
            logits = model(input_tensor)
            logits = F.softmax(logits / request.temperature, dim=-1)
            torch.cuda.synchronize()
        elapsed_ms = (time.perf_counter() - start) * 1000
        
        return PredictResponse(
            logits=logits[0].tolist(),
            inference_time_ms=round(elapsed_ms, 3),
            model_version=model_version
        )

# ---- Graceful Shutdown ----
def handle_sigterm(signum, frame):
    logger.warning("Received SIGTERM. Initiating graceful shutdown.")
    sys.exit(0)

signal.signal(signal.SIGTERM, handle_sigterm)

Multi-GPU and Distributed Training in Containers

For distributed training workloads, containers need inter-container communication via NCCL. Docker Compose can simulate this on a single multi-GPU node, but for production across multiple nodes you'll want Kubernetes with the NVIDIA GPU Operator. Here's a Docker Compose snippet that launches two training workers sharing GPUs via NCCL:

services:
  trainer-worker-0:
    image: myregistry/training:latest
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 2
              capabilities: [gpu]
    environment:
      - CUDA_VISIBLE_DEVICES=0,1
      - MASTER_ADDR=trainer-worker-0
      - MASTER_PORT=29500
      - WORLD_SIZE=2
      - RANK=0
      - NCCL_DEBUG=INFO
      - NCCL_SOCKET_IFNAME=eth0
    command: ["torchrun", "--nproc_per_node=2", "train.py"]
    
  trainer-worker-1:
    image: myregistry/training:latest
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 2
              capabilities: [gpu]
    environment:
      - CUDA_VISIBLE_DEVICES=0,1
      - MASTER_ADDR=trainer-worker-0
      - MASTER_PORT=29500
      - WORLD_SIZE=2
      - RANK=1
      - NCCL_DEBUG=INFO
      - NCCL_SOCKET_IFNAME=eth0
    command: ["torchrun", "--nproc_per_node=2", "train.py"]
    depends_on:
      trainer-worker-0:
        condition: service_healthy

In a real Kubernetes production cluster, you'd replace this with a torchrun launcher inside a single pod requesting multiple GPUs, or use the Kubeflow Training Operator for PyTorchJob resources that handle worker pod orchestration automatically.

Monitoring GPU Metrics in Production

Observability is non-negotiable. Beyond the DCGM exporter shown earlier, instrument your application code with GPU metric emission:

import torch
from prometheus_client import Gauge, CollectorRegistry, generate_latest
import threading
import time

class GPUMetricsCollector:
    def __init__(self, collection_interval: int = 15):
        self.registry = CollectorRegistry()
        self.gpu_utilization = Gauge(
            "gpu_utilization_percent",
            "GPU utilization percentage",
            ["device"],
            registry=self.registry
        )
        self.gpu_memory_used = Gauge(
            "gpu_memory_used_bytes",
            "GPU memory used in bytes",
            ["device"],
            registry=self.registry
        )
        self.gpu_temperature = Gauge(
            "gpu_temperature_celsius",
            "GPU temperature in Celsius",
            ["device"],
            registry=self.registry
        )
        self.inference_latency = Gauge(
            "inference_latency_milliseconds",
            "Model inference latency",
            ["model_version"],
            registry=self.registry
        )
        
        self._stop_event = threading.Event()
        self._collector_thread = threading.Thread(
            target=self._collect_loop,
            args=(collection_interval,),
            daemon=True
        )
    
    def start(self):
        self._collector_thread.start()
    
    def stop(self):
        self._stop_event.set()
        self._collector_thread.join(timeout=5)
    
    def _collect_loop(self, interval: int):
        while not self._stop_event.is_set():
            try:
                for i in range(torch.cuda.device_count()):
                    device_name = f"cuda:{i}"
                    # Query GPU stats via pynvml or torch internal API
                    free_mem, total_mem = torch.cuda.mem_get_info(i)
                    used_mem = total_mem - free_mem
                    self.gpu_memory_used.labels(device=device_name).set(used_mem)
                    
                    # Utilization requires pynvml for accurate readings
                    try:
                        import pynvml
                        pynvml.nvmlInit()
                        handle = pynvml.nvmlDeviceGetHandleByIndex(i)
                        util = pynvml.nvmlDeviceGetUtilizationRates(handle)
                        temp = pynvml.nvmlDeviceGetTemperature(handle, pynvml.NVML_TEMPERATURE_GPU)
                        self.gpu_utilization.labels(device=device_name).set(util.gpu)
                        self.gpu_temperature.labels(device=device_name).set(temp)
                    except Exception:
                        pass
            except Exception as e:
                print(f"GPU metric collection error: {e}")
            self._stop_event.wait(interval)
    
    def get_metrics_bytes(self) -> bytes:
        return generate_latest(self.registry)

Expose these metrics on a separate port (e.g., 9090) and scrape them with Prometheus. Pair this with Grafana dashboards showing GPU utilization, memory pressure, inference latency percentiles, and temperature trends. Alert on GPU memory usage exceeding 90% or temperature crossing 80Ā°C.

Best Practices for Production GPU Containers

1. Pin Exact Versions Everywhere

GPU ecosystems are fragile. Pin the CUDA base image digest (not just tag), the cuDNN version, the PyTorch patch number, and every pip dependency with hash checking. A single floating dependency can pull an incompatible cuDNN build that fails silently on certain GPU architectures.

# Example: pinning with sha256 digest
FROM nvidia/cuda:12.4.1-devel-ubuntu22.04@sha256:abc123def456...
RUN pip install --require-hashes -r requirements.txt

2. Use Multi-Stage Builds Religiously

Never ship CUDA development tools (nvcc, cuda-gdb, headers) in your runtime image. They add 2-3 GB of bloat and are useless at inference time. The builder/runtime pattern shown earlier is the gold standard.

3. Warm Up GPUs Before Serving Traffic

CUDA kernels compile lazily. The first inference request can take 3-10 seconds due to JIT compilation. Send a dummy batch through the model during container startup to force kernel compilation and cache the compiled kernels.

4. Implement Proper Graceful Shutdown

GPU containers need time to drain in-flight requests and free CUDA memory. Kubernetes sends SIGTERM and waits for terminationGracePeriodSeconds (default 30s). Your application must catch SIGTERM, stop accepting new requests, finish existing ones, and call torch.cuda.empty_cache() before exiting.

5. Separate Model Weights from Container Images

Don't bake 14 GB model checkpoints into container imagesā€”it makes image pulls excruciatingly slow. Mount models via volumes, init containers, or lazy-load from object storage. Reserve image-stored models only for tiny ensembles (under 100 MB).

6. Run as Non-Root User

CUDA containers default to root, which is a security risk. Create an appuser with minimal permissions as shown in the Dockerfile. The NVIDIA container toolkit handles device permissions correctly for non-root users.

7. Configure NCCL Environment Variables for Scale

When running multi-GPU or multi-node, tune these environment variables:

ENV NCCL_DEBUG=INFO \
    NCCL_SOCKET_IFNAME=eth0 \
    NCCL_IB_DISABLE=0 \
    NCCL_NET_GDR_LEVEL=2 \
    NCCL_MIN_NCHANNELS=16 \
    OMPI_MCA_btl_vader_single_copy_mechanism=none

These settings enable InfiniBand GPU Direct RDMA, disable single-copy mechanisms that conflict with CUDA, and increase channel counts for better all-reduce throughput.

8. Automate Image Building in CI/CD

GPU Docker images take longer to build due to large CUDA layers. Cache aggressively in your CI system. Use Docker BuildKit with --cache-from and remote cache stores. Run GPU integration tests on every imageā€”a container that builds successfully can still crash at runtime due to driver-kernel mismatches.

Common Pitfalls and How to Avoid Them

Conclusion

Dockerized GPU ML workloads bring the full power of containerizationā€”reproducibility, portability, density, and operational simplicityā€”to the most demanding compute workloads in modern infrastructure. By combining NVIDIA's CUDA base images with multi-stage builds, explicit GPU resource requests, proper health checking, and comprehensive metrics collection, you create a production system that serves inferences reliably, scales training jobs across nodes, and integrates seamlessly with Kubernetes or Docker Compose-based deployments. The practices outlined hereā€”pinned versions, non-root users, warm-up passes, graceful shutdown, and tiered model cachingā€”form the foundation of any mature GPU ML platform. Start with a single well-crafted Dockerfile, add the monitoring and resource controls incrementally, and you'll have a GPU service that's as manageable and predictable as any other containerized workload in your fleet.

šŸš€ Need a reliable AI agent for your project?

Deploy Hermes Agent in 10 minutes. Managed hosting, zero DevOps.

Get Started ā€” $23.99/mo
ā† Back to all articles