Welcome to from-docker-to-kubernetes

Docker for AI/ML Workloads

Learn how to effectively containerize, deploy, and orchestrate AI and machine learning workloads with Docker

H A R S H H A A
@NotHarshhaa

Docker for AI/ML Workloads

Docker provides an excellent platform for developing, training, and deploying AI and machine learning models, offering reproducibility, portability, and scalability for complex AI workflows. By containerizing AI/ML environments, data scientists and ML engineers can ensure consistent execution across development, testing, and production systems while eliminating the notorious "it works on my machine" problem that often plagues complex ML dependencies. Docker also enables efficient collaboration between teams, version control of entire environments, and seamless integration with orchestration tools for distributed training and inference.

AI/ML Development Environment

Base Images for AI/ML

  • NVIDIA CUDA images for GPU workloads
    • Pre-configured with CUDA drivers and libraries
    • Various versions available to match specific CUDA requirements
    • Optimized for different GPU architectures (Pascal, Volta, Turing, Ampere)
    • Base layer for building custom deep learning environments
    • Example: nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04
  • TensorFlow official images
    • Complete environments with TensorFlow pre-installed
    • Available in CPU and GPU variants
    • Jupyter notebooks integration in many images
    • Consistent versioning with TensorFlow releases
    • Example: tensorflow/tensorflow:2.12.0-gpu
  • PyTorch container ecosystem
    • Official PyTorch installations with CUDA support
    • Optimized for performance with GPU acceleration
    • Includes common PyTorch libraries and extensions
    • Various Python version options
    • Example: pytorch/pytorch:2.0.1-cuda11.7-cudnn8-runtime
  • Scikit-learn and data science stacks
    • Comprehensive Python data science toolkits
    • Pandas, NumPy, Matplotlib, and other common libraries
    • Ready-to-use Jupyter environments
    • Optimized for data processing pipelines
    • Example: jupyter/scipy-notebook:python-3.10
  • Specialized deep learning images
    • Domain-specific containers (NLP, computer vision, etc.)
    • Pre-trained models and frameworks
    • Hugging Face Transformers, Detectron2, etc.
    • Production-optimized inference containers
    • Example: huggingface/transformers-pytorch-gpu:4.29.2

Setting Up GPU Support

  • NVIDIA Container Toolkit (nvidia-docker)
    • System component that enables GPU access from containers
    • Provides runtime extensions for Docker
    • Manages NVIDIA driver mapping between host and container
    • Essential prerequisite for GPU-accelerated containers
    • Installation: apt-get install nvidia-container-toolkit
  • GPU passthrough configuration
    • Exposes specific GPUs to containers
    • Controls which containers can access which GPUs
    • Enables fine-grained resource allocation
    • Configured with --gpus flag or in docker-compose
    • Example: docker run --gpus device=0,1 nvidia/cuda nvidia-smi
  • Driver compatibility considerations
    • Container CUDA version must be compatible with host driver
    • Driver version must support the required CUDA version
    • Compatibility matrix available in NVIDIA documentation
    • Minimum driver version depends on CUDA toolkit version
    • Best practice: Use container CUDA version ≤ host driver CUDA capability
  • Resource allocation
    • Memory limits for preventing OOM errors
    • GPU memory monitoring and management
    • NVIDIA MPS for shared GPU access
    • Using nvidia-smi for resource monitoring
    • Example: nvidia-smi --query-gpu=utilization.gpu,memory.used --format=csv
  • Multi-GPU setups
    • Data parallelism across multiple GPUs
    • Managing process-to-GPU mapping
    • NCCL configuration for inter-GPU communication
    • NVLink considerations for high-bandwidth connections
    • Optimizing container placement for GPU topology

Essential Docker Images for AI/ML

# TensorFlow with GPU support
FROM tensorflow/tensorflow:latest-gpu
# This image includes:
# - TensorFlow framework with GPU acceleration
# - Python and essential scientific libraries
# - Pre-configured CUDA and cuDNN
# - Compatible with NVIDIA Container Toolkit
# - Jupyter server for interactive development

# PyTorch with CUDA
FROM pytorch/pytorch:latest
# This image includes:
# - PyTorch deep learning framework
# - CUDA and cuDNN optimized for PyTorch
# - TorchVision, TorchAudio, and TorchText
# - Python 3 with scientific computing packages
# - Optimized performance for training and inference

# RAPIDS for GPU-accelerated data science
FROM rapidsai/rapidsai:latest
# This image includes:
# - RAPIDS suite for GPU-accelerated data science
# - cuDF (GPU DataFrame library similar to pandas)
# - cuML (Machine learning algorithms on GPU)
# - cuGraph (Graph analytics on GPU)
# - Dask for distributed computing
# - Integration with scikit-learn API

# Scikit-learn and common data science tools
FROM jupyter/scipy-notebook:latest
# This image includes:
# - Jupyter Lab/Notebook server
# - pandas, NumPy, Matplotlib, Seaborn
# - scikit-learn, SciPy, and StatsModels
# - Patsy and other data analysis tools
# - Comprehensive scientific Python stack
# - Ready for CPU-based machine learning

# Hugging Face Transformers
FROM huggingface/transformers-pytorch-gpu:latest
# This image includes:
# - Transformers library for NLP tasks
# - Pre-trained models and tokenizers
# - PyTorch with GPU acceleration
# - Optimized for inference and fine-tuning
# - Support for BERT, GPT, T5, and other architectures

Optimizing Dockerfiles for ML

# Example Dockerfile for ML development
FROM python:3.10-slim as builder

WORKDIR /app
RUN apt-get update && apt-get install -y --no-install-recommends \
    build-essential \
    && rm -rf /var/lib/apt/lists/*

COPY requirements.txt .
RUN pip wheel --no-cache-dir --no-deps --wheel-dir /app/wheels -r requirements.txt

FROM python:3.10-slim

WORKDIR /app
COPY --from=builder /app/wheels /wheels
RUN pip install --no-cache /wheels/*

COPY src/ ./src
COPY models/ ./models
COPY config.yaml .

ENV PYTHONUNBUFFERED=1 \
    PYTHONDONTWRITEBYTECODE=1 \
    MODEL_PATH=/app/models \
    WORKERS=4

ENTRYPOINT ["python", "src/serve.py"]

Data Management Strategies

Volume Mounting for Datasets

  • Mount large datasets as volumes
    • Avoids copying data into containers
    • Enables sharing datasets across multiple containers
    • Persists data beyond container lifecycle
    • Improves build time and reduces image size
    • Example: docker run -v /host/data:/data:ro tensorflow/tensorflow:latest-gpu
  • Configure data caching
    • Implement multi-level caching strategies
    • Use tmpfs mounts for high-speed temporary storage
    • Leverage SSD for intermediate datasets and HDD for archival
    • Cache preprocessed data to avoid redundant computation
    • Example: --mount type=tmpfs,destination=/cache,tmpfs-size=16g
  • Structure data directories
    • Organize by dataset, version, and splits (train/val/test)
    • Implement consistent naming conventions
    • Create metadata files documenting dataset properties
    • Consider columnar formats (Parquet, Arrow) for efficiency
    • Design for parallel access patterns in distributed training
  • Version datasets effectively
    • Implement content-addressable storage patterns
    • Use dataset versioning tools (DVC, Pachyderm, etc.)
    • Create dataset manifests with checksums
    • Track dataset lineage and transformations
    • Consider ACID-compliant dataset management
  • Optimize I/O operations
    • Use memory mapping for large files
    • Implement asynchronous data loading pipelines
    • Consider data compression tradeoffs
    • Tune buffer sizes for specific storage systems
    • Example: TensorFlow tf.data API with prefetching and parallelism

Model Storage and Versioning

  • Efficient model serialization
    • Choose appropriate serialization formats (saved_model, ONNX, TorchScript)
    • Optimize for size vs. loading speed tradeoffs
    • Consider quantization for model compression
    • Implement selective parameter saving
    • Example: Convert large models to fp16 precision for storage
  • Model registry integration
    • Connect to MLflow, Weights & Biases, or custom registries
    • Implement automated versioning on training completion
    • Tag models with metadata (metrics, dataset version, etc.)
    • Support model lineage tracking
    • Example: mlflow.tensorflow.log_model(model, "model")
  • Version control for models
    • Store models with semantic versioning
    • Implement immutable model artifacts
    • Track model dependencies and environment
    • Manage experimental vs. production models
    • Example: Use Git LFS or specialized model versioning tools
  • Artifact management
    • Define lifecycle policies for model retention
    • Implement access control for model artifacts
    • Store evaluation metrics alongside models
    • Include sample inputs/outputs with models
    • Track compute resources used for training
  • Reproducible model loading
    • Store model configuration separately from weights
    • Document initialization procedures
    • Version model loaders alongside models
    • Implement model compatibility checking
    • Example: Create model cards with reproducibility instructions

Training Workflows

# docker-compose.yml for ML training
version: '3.8'
services:
  training:
    build: 
      context: .
      dockerfile: Dockerfile.train
    volumes:
      - ./data:/data
      - ./output:/output
    environment:
      - EPOCHS=100
      - BATCH_SIZE=32
      - LEARNING_RATE=0.001
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [gpu]

Inference Serving

Model Serving Options

  • TensorFlow Serving
    • Production-grade serving system for TensorFlow models
    • Supports model versioning and hot swapping
    • Highly optimized for TensorFlow SavedModel format
    • Provides both gRPC and REST APIs
    • Enables batching and high-performance inference
    • Example: tensorflow/serving:2.12.0
  • NVIDIA Triton Inference Server
    • Multi-framework inference server (TensorFlow, PyTorch, ONNX, etc.)
    • Dynamic batching and sequence batching
    • Concurrent model execution
    • Model ensemble support
    • Optimized for NVIDIA GPUs with TensorRT integration
    • Example: nvcr.io/nvidia/tritonserver:23.04-py3
  • TorchServe
    • Production serving system for PyTorch models
    • Model versioning and management
    • REST and gRPC endpoints
    • A/B testing capabilities
    • Custom handlers for preprocessing/postprocessing
    • Example: pytorch/torchserve:0.7.1-gpu
  • ONNX Runtime
    • Cross-platform inference engine for ONNX models
    • Hardware acceleration on various devices (CPU, GPU, TPU)
    • Quantization and optimization support
    • Wide framework compatibility
    • Graph optimizations for performance
    • Example: mcr.microsoft.com/azureml/onnxruntime:latest
  • Custom REST API services
    • Flexible APIs built on Flask, FastAPI, or other frameworks
    • Complete control over request handling and processing
    • Easy integration with business logic
    • Custom authentication and authorization
    • Tailored scaling and deployment options
    • Example: FastAPI with model loading on startup

Optimizing Inference Containers

# Optimized TensorFlow Serving container
FROM tensorflow/serving:latest

# Copy the SavedModel
COPY ./models/saved_model /models/my_model/1
# Version directory structure is important for TF Serving
# /models/my_model/1 indicates version 1 of the model

# Set environment variables
ENV MODEL_NAME=my_model
# This sets the base model name that will be served

# Expose ports for REST and gRPC
EXPOSE 8501 8500
# 8501: RESTful API port for HTTP requests
# 8500: gRPC port for high-performance clients

# Configure optimizations
ENV TF_CPP_MIN_LOG_LEVEL=2 \
    TF_ENABLE_ONEDNN_OPTS=1 \
    OMP_NUM_THREADS=4 \
    MALLOC_TRIM_THRESHOLD_=0
# TF_CPP_MIN_LOG_LEVEL=2: Suppress info and warning logs
# TF_ENABLE_ONEDNN_OPTS=1: Enable Intel MKL-DNN optimizations
# OMP_NUM_THREADS=4: Control thread parallelism
# MALLOC_TRIM_THRESHOLD_=0: Disable memory trimming for performance

# Set entrypoint
ENTRYPOINT ["tensorflow_model_server", "--port=8500", "--rest_api_port=8501", "--model_config_file=/models/models.config"]
# models.config allows serving multiple models from the same server
# Configuration includes model name, platform, and version policy

Distributed Training

# Example of distributed training with Docker Compose
version: '3.8'
services:
  parameter-server:
    image: my-ml-training:latest
    command: ["python", "distributed_train.py", "--job_name=ps", "--task_index=0"]
    # Parameter server coordinates distributed training
    # Manages model parameters and optimization
    # Aggregates gradients from workers
    # Distributes updated parameters back to workers
    ports:
      - "2222:2222"  # Port for worker communication
    volumes:
      - ./data:/data  # Mount dataset volume
      - ./output:/output  # Mount for saving results and checkpoints
    environment:
      - TF_CONFIG={"cluster":{"ps":["parameter-server:2222"],"worker":["worker-0:2223","worker-1:2223"]},"task":{"type":"ps","index":0}}
    networks:
      - training-network  # Dedicated network for training communication
    deploy:
      resources:
        limits:
          cpus: '4.0'
          memory: 8G  # Parameter servers are often CPU/RAM intensive
      restart_policy:
        condition: on-failure
        max_attempts: 3

  worker-0:
    image: my-ml-training:latest
    command: ["python", "distributed_train.py", "--job_name=worker", "--task_index=0"]
    # Worker 0 performs training computation
    # Processes data batches and computes gradients
    # Communicates with parameter server
    # Can handle part of the dataset or model
    depends_on:
      - parameter-server  # Ensures parameter server is started first
    ports:
      - "2223:2223"  # Worker communication port
    volumes:
      - ./data:/data  # Read-only access to training data
    environment:
      - TF_CONFIG={"cluster":{"ps":["parameter-server:2222"],"worker":["worker-0:2223","worker-1:2223"]},"task":{"type":"worker","index":0}}
    networks:
      - training-network
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1  # Reserves one GPU
              capabilities: [gpu]  # Enables GPU capabilities
        limits:
          cpus: '8.0'
          memory: 16G  # Workers need sufficient memory for batches
      restart_policy:
        condition: on-failure

  worker-1:
    image: my-ml-training:latest
    command: ["python", "distributed_train.py", "--job_name=worker", "--task_index=1"]
    # Worker 1 processes a different portion of data
    # Operates in parallel with worker-0
    # Provides additional computing power
    # Enables data parallelism across multiple GPUs
    depends_on:
      - parameter-server
    ports:
      - "2224:2223"  # Different host port to avoid conflicts
    volumes:
      - ./data:/data
    environment:
      - TF_CONFIG={"cluster":{"ps":["parameter-server:2222"],"worker":["worker-0:2223","worker-1:2223"]},"task":{"type":"worker","index":1}}
    networks:
      - training-network
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [gpu]
        limits:
          cpus: '8.0'
          memory: 16G
      restart_policy:
        condition: on-failure

networks:
  training-network:
    driver: bridge
    # Dedicated network optimizes inter-container communication
    # Isolates training traffic from other services
    # Can be configured for high-throughput, low-latency communication

MLOps Integration

Orchestrating ML Pipelines

Pipeline Components

  • Data preparation containers
    • Data validation and cleaning
    • Format conversion and normalization
    • Feature extraction from raw data
    • Dataset splitting (train/validation/test)
    • Data augmentation for training
    • Example: Apache Beam or Luigi containers
  • Feature engineering services
    • Feature transformation pipelines
    • Feature selection algorithms
    • Dimensionality reduction
    • Feature encoding and normalization
    • Feature store integration
    • Example: Feast or custom feature services
  • Model training jobs
    • Hyperparameter optimization
    • Model fitting and validation
    • Cross-validation execution
    • Checkpoint management
    • Distributed training coordination
    • Example: Containers with TensorFlow, PyTorch, etc.
  • Evaluation workers
    • Model performance assessment
    • Metric calculation and validation
    • A/B comparison with baseline models
    • Threshold determination
    • Test dataset evaluation
    • Example: Custom containers running evaluation scripts
  • Deployment services
    • Model packaging for production
    • Serving infrastructure setup
    • Canary deployment handling
    • Versioning and rollback support
    • Integration with API gateways
    • Example: KServe or TensorFlow Serving containers
  • Monitoring components
    • Data drift detection
    • Model performance tracking
    • Resource utilization monitoring
    • Prediction logging and analysis
    • Alert generation for degradation
    • Example: Prometheus exporters and Grafana dashboards

Workflow Orchestration

# Example Argo Workflow for ML pipeline
apiVersion: argoproj.io/v1alpha1
kind: Workflow
metadata:
  name: ml-training-pipeline
  # Defines a named ML training pipeline
  # Will appear in Argo UI with this identifier
  # Can be triggered manually or by events
spec:
  entrypoint: ml-pipeline
  # Main entry point for the workflow
  # Defines where execution begins
  
  # Optional workflow-wide settings
  # ttlStrategy defines how long to keep workflow after completion
  ttlStrategy:
    secondsAfterCompletion: 86400  # 24 hours
  
  # Optional volume claims for persistent storage
  volumeClaimTemplates:
  - metadata:
      name: workdir
    spec:
      accessModes: [ "ReadWriteOnce" ]
      resources:
        requests:
          storage: 10Gi
  
  templates:
  - name: ml-pipeline
    # Main pipeline template that orchestrates all steps
    steps:
    - - name: data-preparation
        # First step: prepare and validate data
        # Runs data cleaning, normalization, and splitting
        # Produces a validated dataset for training
        template: data-prep
        # Detailed implementation of data-prep is defined elsewhere
        
    - - name: model-training
        # Second step: train the ML model
        # Uses the prepared data to fit model parameters
        # Outputs trained model artifacts
        template: train
        arguments:
          parameters:
          - name: data-path
            # Dynamic parameter from previous step
            # Allows passing the location of prepared data
            value: "{{steps.data-preparation.outputs.parameters.data-path}}"
        
    - - name: model-evaluation
        # Third step: evaluate model performance
        # Calculates metrics on validation data
        # Determines if model meets quality thresholds
        template: evaluate
        arguments:
          parameters:
          - name: model-path
            # Reference to the trained model from previous step
            value: "{{steps.model-training.outputs.parameters.model-path}}"
        
    - - name: model-deployment
        # Fourth step: deploy model to production
        # Only executes if evaluation score is sufficient
        # Handles model serving infrastructure
        template: deploy
        arguments:
          parameters:
          - name: model-path
            # Location of the model to deploy
            value: "{{steps.model-training.outputs.parameters.model-path}}"
          - name: evaluation-result
            # Evaluation metric to include in deployment metadata
            value: "{{steps.model-evaluation.outputs.parameters.result}}"
        # Conditional execution based on model quality
        # Only deploys if accuracy exceeds 85%
        when: "{{steps.model-evaluation.outputs.parameters.result}} > 0.85"
        
  # Additional templates would define the implementation details
  # of data-prep, train, evaluate, and deploy steps
  - name: data-prep
    container:
      image: my-registry/data-processor:v1
      command: [python, data_prep.py]
      # Implementation details...
      
  - name: train
    container:
      image: my-registry/model-trainer:v1
      # Implementation details...
      
  - name: evaluate
    container:
      image: my-registry/model-evaluator:v1
      # Implementation details...
      
  - name: deploy
    container:
      image: my-registry/model-deployer:v1
      # Implementation details...

Resource Management

# GPU allocation in Docker
docker run --gpus all -it tensorflow/tensorflow:latest-gpu python -c "import tensorflow as tf; print(tf.config.list_physical_devices('GPU'))"
# Lists all available GPUs inside the container
# Verifies that TensorFlow can see the GPUs
# Confirms proper NVIDIA Container Toolkit setup
# Useful for debugging GPU visibility issues

# Specific GPU selection
docker run --gpus '"device=1,2"' -it pytorch/pytorch:latest python -c "import torch; print(torch.cuda.device_count())"
# Selects specific GPUs by index (devices 1 and 2)
# Useful for multi-tenant environments
# Allows fine-grained resource allocation
# Prevents container from using all available GPUs
# Alternative syntax: --gpus '"device=1,capabilities=compute,utility"'

# GPU memory limits (with NVIDIA MPS)
docker run --gpus all --env NVIDIA_MPS_ACTIVE=1 --env CUDA_MPS_ACTIVE_THREAD_PERCENTAGE=20 -it tensorflow/tensorflow:latest-gpu
# Enables NVIDIA Multi-Process Service (MPS)
# Allows multiple containers to share GPU resources more efficiently
# Controls percentage of GPU compute resources allocated
# Improves GPU utilization for smaller workloads
# Useful for serving multiple models on the same GPU


## Experiment Tracking

::steps
### Container-based Experiment Management
- MLflow containers for tracking
  - Open-source platform for ML lifecycle management
  - Experiment tracking and comparison
  - Model registry and versioning
  - Centralized metrics and artifacts storage
  - REST API for programmatic access to results
  - Example: `ghcr.io/mlflow/mlflow:latest`
  
- Weights & Biases integration
  - Cloud-based experiment tracking service
  - Real-time training visualization
  - Hyperparameter importance analysis
  - Collaborative experiment management
  - Model and dataset versioning
  - Example: `wandb/local` for self-hosted option
  
- TensorBoard deployment
  - TensorFlow's visualization toolkit
  - Training metrics visualization
  - Graph visualization for neural networks
  - Embedding projections and feature analysis
  - Model profiling and debugging tools
  - Example: `tensorflow/tensorflow:latest` includes TensorBoard
  
- Custom metrics collection
  - Specialized metric collection APIs
  - Performance counters for hardware utilization
  - Domain-specific evaluation metrics
  - A/B testing frameworks
  - Real-time alerting on metric thresholds
  - Example: Custom Flask API containers for metrics
  
- Experiment versioning
  - Git integration for code versioning
  - Environment snapshots for reproducibility
  - Configuration management (with Hydra or similar)
  - Parameter versioning and comparison
  - Experimental lineage tracking
  - Example: DVC for data version control with experiments

### Example Setup
```yaml
version: '3.8'
services:
  mlflow:
    image: ghcr.io/mlflow/mlflow:latest
    ports:
      - "5000:5000"    # Web UI and API port
    volumes:
      - ./mlflow:/mlflow    # Persistent storage for experiment data
    command: ["mlflow", "server", "--host", "0.0.0.0", "--backend-store-uri", "sqlite:///mlflow/mlflow.db", "--default-artifact-root", "/mlflow/artifacts"]
    # Uses SQLite database for metadata storage
    # Configures local file system for artifact storage
    # Can be scaled with external databases like PostgreSQL
    # Host 0.0.0.0 allows external connections
    environment:
      - MLFLOW_S3_ENDPOINT_URL=http://minio:9000    # Optional object storage
      - AWS_ACCESS_KEY_ID=minioadmin              # For S3-compatible storage
      - AWS_SECRET_ACCESS_KEY=minioadmin          # For S3-compatible storage
    networks:
      - ml-network
    restart: unless-stopped
    healthcheck:
      test: ["CMD", "curl", "-f", "http://localhost:5000/api/2.0/mlflow/experiments/list"]
      interval: 30s
      timeout: 10s
      retries: 3
  
  minio:
    image: minio/minio:latest
    ports:
      - "9000:9000"    # API port
      - "9001:9001"    # Console port
    volumes:
      - ./minio-data:/data
    command: server /data --console-address ":9001"
    environment:
      - MINIO_ROOT_USER=minioadmin
      - MINIO_ROOT_PASSWORD=minioadmin
    networks:
      - ml-network
  
  notebook:
    build: ./notebooks
    ports:
      - "8888:8888"    # Jupyter notebook interface
    volumes:
      - ./notebooks:/home/jovyan/work    # Notebook source files
      - ./data:/home/jovyan/data         # Dataset access
    environment:
      - MLFLOW_TRACKING_URI=http://mlflow:5000    # Connect to MLflow
      - PYTHONPATH=/home/jovyan/work              # For importing local modules
      - WANDB_API_KEY=${WANDB_API_KEY:-}          # Optional W&B integration
      - JUPYTER_ENABLE_LAB=yes                    # Enable JupyterLab interface
    depends_on:
      - mlflow
    networks:
      - ml-network
    restart: unless-stopped
    command: start-notebook.sh --NotebookApp.token='' --NotebookApp.password=''

networks:
  ml-network:
    driver: bridge

::

Hyperparameter Optimization

Deployment Architectures

Edge Deployment

  • Optimized containers for edge devices
    • Minimized container size for limited storage
    • Platform-specific builds (ARM, x86, RISC-V)
    • Specialized base images (Alpine, Distroless)
    • Static linking to reduce dependencies
    • Example: FROM arm32v7/python:3.9-slim for Raspberry Pi deployment
  • Model quantization and pruning
    • Int8/FP16 quantization for reduced memory footprint
    • Weight pruning for smaller model size
    • Knowledge distillation for compact student models
    • Post-training optimization techniques
    • Example: TensorFlow Lite models with 75% size reduction
  • Runtime optimization
    • Hardware-specific acceleration (NEON, AVX)
    • Memory mapping for efficient loading
    • Thread and process optimization
    • Batch size tuning for latency vs throughput
    • Example: ONNX Runtime with custom execution providers
  • Resource-constrained environments
    • CPU/RAM/storage limitations management
    • Thermal and power consumption considerations
    • Offline operation capabilities
    • Graceful degradation under resource pressure
    • Example: Container configured with --memory=512m --cpus=0.5
  • Update strategies for edge models
    • Delta updates to minimize bandwidth
    • A/B model deployment for validation
    • Rollback mechanisms for failed updates
    • Version compatibility verification
    • Example: Container image layering for efficient updates

Cloud Deployment

  • Scalable inference APIs
    • RESTful and gRPC API interfaces
    • Stateless design for horizontal scaling
    • Asynchronous processing for batch requests
    • Client libraries for multiple languages
    • Example: KServe or TorchServe behind API Gateway
  • Auto-scaling model servers
    • Horizontal pod autoscaling based on CPU/memory/custom metrics
    • Prediction request queue-based scaling
    • Minimum replicas for baseline performance
    • GPU utilization-based scaling policies
    • Example: Kubernetes HPA with custom metrics from Prometheus
  • Load balancing strategies
    • Round-robin for stateless inference
    • Session affinity for stateful models
    • Weighted distribution based on instance capacity
    • Latency-based routing for global deployments
    • Example: Cloud load balancer with health checks
  • High-availability configurations
    • Multi-zone and multi-region deployments
    • Automated failover mechanisms
    • Redundant model server instances
    • State replication where needed
    • Example: Multi-regional Kubernetes clusters with PodDisruptionBudget
  • Cloud-native integrations
    • Managed Kubernetes services (EKS, GKE, AKS)
    • Serverless inference (AWS Lambda, Cloud Run, Azure Functions)
    • Cloud monitoring and logging integration
    • Identity and access management integration
    • Example: AWS SageMaker with auto-scaling inference endpoints

Performance Optimization

# Performance-optimized Dockerfile for inference
FROM python:3.10-slim

# Install performance libraries
RUN apt-get update && apt-get install -y --no-install-recommends \
    libopenblas-dev \
    libomp-dev \
    && rm -rf /var/lib/apt/lists/*
# libopenblas-dev: Optimized BLAS implementation for linear algebra operations
# libomp-dev: OpenMP runtime for parallel processing
# Cleaning apt cache reduces image size

# Install optimized packages
RUN pip install --no-cache-dir \
    numpy==1.24.* \
    onnxruntime-gpu==1.15.* \
    onnx==1.14.* \
    optimum==1.11.*
# numpy: Pinned version for stability and compatibility
# onnxruntime-gpu: Hardware-accelerated inference engine
# onnx: Open Neural Network Exchange format support
# optimum: Hugging Face's optimization toolkit
# --no-cache-dir reduces image size

# Copy model and application
COPY ./model /app/model
COPY ./src /app/src
WORKDIR /app
# Separate model and code copying allows for better layer caching
# Models change less frequently than code in many scenarios

# Set optimization environment variables
ENV OMP_NUM_THREADS=4 \
    OMP_WAIT_POLICY=ACTIVE \
    OPENBLAS_NUM_THREADS=4 \
    ONNXRUNTIME_CUDA_DEVICE_ID=0
# OMP_NUM_THREADS: Controls thread parallelism for OpenMP
# OMP_WAIT_POLICY=ACTIVE: Keeps threads active for faster response
# OPENBLAS_NUM_THREADS: Controls threading in linear algebra operations
# ONNXRUNTIME_CUDA_DEVICE_ID: Selects specific GPU for inference

# Add performance monitoring capabilities
RUN pip install --no-cache-dir prometheus_client==0.17.* py-spy==0.3.*
# prometheus_client: Exposes metrics for monitoring
# py-spy: Low-overhead profiling for Python processes

# Configure memory optimizations
ENV MALLOC_TRIM_THRESHOLD_=65536 \
    PYTHONMALLOC=malloc \
    PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128
# MALLOC_TRIM_THRESHOLD_: Controls memory deallocation behavior
# PYTHONMALLOC=malloc: Uses system allocator instead of Python's
# PYTORCH_CUDA_ALLOC_CONF: Optimizes GPU memory fragmentation

# Run optimized server with monitoring
CMD ["python", "src/serve_optimized.py"]
# Entry point runs the optimized serving code
# Consider using gunicorn or uvicorn for production HTTP servers

Real-world ML Use Cases

Advanced Topics

Multi-node Training

  • Container orchestration for distributed training
    • Kubernetes for managing training pods across nodes
    • Custom operators for ML workloads (KubeFlow, Ray)
    • Resource allocation optimization for heterogeneous clusters
    • Training job scheduling and prioritization
    • Example: Kubernetes StatefulSets for ordered pod creation
  • Parameter servers and workers
    • Architectural patterns for distributed optimization
    • Sharded parameter servers for large models
    • Asynchronous vs. synchronous parameter updates
    • Communication topology optimization
    • Example: TensorFlow tf.distribute.ParameterServerStrategy
  • Network optimization for data transfer
    • RDMA/RoCE for high-speed GPU communication
    • Gradient compression techniques to reduce bandwidth
    • Topology-aware pod placement
    • Custom container networking plugins
    • Example: NVIDIA NCCL with InfiniBand for GPU communication
  • Checkpoint management
    • Distributed checkpoint coordination
    • Incremental checkpointing strategies
    • Cloud storage integration for durability
    • Checkpoint validation and corruption detection
    • Example: Distributed TensorFlow CheckpointManager
  • Failure recovery strategies
    • Preemption-aware training processes
    • Automatic worker replacement
    • Elastic training group management
    • Gradual scaling with minimal recomputation
    • Example: PyTorch Elastic for fault-tolerant training

Federated Learning

  • Container-based federated learning nodes
    • Self-contained training environments on edge devices
    • Minimal runtime dependencies for diverse deployments
    • Standardized APIs for model and update exchange
    • Resource-constrained container optimization
    • Example: TensorFlow Federated client containers
  • Secure aggregation strategies
    • Cryptographic protocols in containerized services
    • Secure multi-party computation containers
    • Zero-knowledge proof systems
    • Threshold cryptography implementations
    • Example: PySyft containers for secure aggregation
  • Privacy-preserving techniques
    • Differential privacy implementation containers
    • Local vs. global privacy budget management
    • Privacy-preserving preprocessing pipelines
    • Anonymization service containers
    • Example: TensorFlow Privacy with configurable DP parameters
  • Edge-to-cloud coordination
    • Asynchronous update mechanisms
    • Connection management for intermittent availability
    • Bandwidth-aware synchronization strategies
    • Multi-tier aggregation hierarchies
    • Example: MQTT-based communication for lightweight coordination
  • Model update synchronization
    • FedAvg and advanced aggregation algorithms
    • Weight divergence monitoring
    • Conflict resolution for concurrent updates
    • Version control for model iterations
    • Example: Flower framework for federated learning orchestration

Best Practices

Troubleshooting Guide

Common Issues

  • GPU not detected in container
    • NVIDIA Container Toolkit not installed or configured properly
    • Incorrect --gpus flag usage or missing GPU capabilities
    • Driver/CUDA version incompatibility
    • GPU visibility issues in nested virtualization
    • Permission problems accessing GPU devices
    • Example error: "could not select device driver with capabilities: [[gpu]]"
  • Out of memory errors during training
    • Batch size too large for available GPU memory
    • Memory leaks from not releasing tensors properly
    • Insufficient container memory limits
    • Fragmented GPU memory after long training
    • Multiple processes competing for same GPU
    • Example error: "CUDA out of memory. Tried to allocate 2.00 GiB"
  • Model loading failures
    • Incompatible serialization format versions
    • Missing model files or incorrect paths
    • Framework version mismatches between save/load
    • Corrupted model files from interrupted saves
    • Insufficient permissions for model directories
    • Example error: "Error loading model: KeyError: 'unexpected key in state_dict'"
  • Performance degradation
    • CPU throttling due to thermal issues
    • Resource contention with other containers
    • Inefficient data loading creating bottlenecks
    • Network saturation in distributed training
    • Suboptimal container resource limits
    • Example symptom: Training iterations becoming progressively slower
  • Data access bottlenecks
    • Inefficient volume mounts or network storage
    • Missing data caching strategies
    • Sequential data access patterns
    • Improper buffer sizes for I/O operations
    • Container networking limitations
    • Example symptom: High wait times in I/O profiling

Diagnostics

# Check GPU visibility
docker run --gpus all nvidia/cuda:11.8.0-base-ubuntu22.04 nvidia-smi
# Shows all available GPUs and their utilization
# Verifies NVIDIA Container Toolkit is working
# Displays driver version and CUDA compatibility
# Shows current GPU memory usage
# Essential first step for GPU troubleshooting

# Debug memory issues
docker stats
# Real-time container resource usage metrics
# Shows CPU, memory, I/O, and network usage
# Helps identify containers approaching resource limits
# Monitor during training to detect memory growth patterns
# Add --no-stream for point-in-time snapshot

# Profile container performance
docker run --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -it my-ml-image:latest
# Additional capabilities for deep profiling
# Allows tools like strace, perf, and py-spy to work
# Enables core dumps for debugging crashes
# Gives visibility into system calls and process behavior
# Example usage inside container: py-spy top --pid 1

# Inspect container logs
docker logs ml-training-container
# Shows stdout/stderr output from the container
# Add --tail=100 to see only recent logs
# Use -f to follow logs in real-time
# Look for error messages and stack traces
# Add --timestamps to correlate with other events

# Interactive debugging
docker exec -it ml-serving-container /bin/bash
# Opens interactive shell in running container
# Allows direct inspection of file system and processes
# Can run diagnostic commands inside container environment
# Access to framework-specific debugging tools
# Example debugging commands inside container:
#   python -c "import torch; print(torch.cuda.is_available())"
#   ps aux | grep python
#   ls -la /app/models
#   cat /proc/1/limits
#   df -h

# GPU profiling with NVIDIA tools
docker run --gpus all -it --rm --pid=host nvidia/cuda:11.8.0-base-ubuntu22.04 nvidia-smi pmon -c 10
# Monitor GPU processes across all containers
# Shows GPU utilization per process
# Identifies which containers are using GPU resources
# Helps detect resource contention issues
# Useful for multi-tenant GPU environments

# Framework-specific debugging
docker exec ml-training-container python -c "import tensorflow as tf; tf.debugging.set_log_device_placement(True); tf.constant(1)"
# Runs diagnostic code within the container
# Shows device placement decisions by framework
# Verifies framework can access appropriate hardware
# Isolates framework-specific configuration issues
# Can be adapted for PyTorch, JAX, or other frameworks
Edit this page
Content Trust & Image Signing
Learn how to implement Docker Content Trust and image signing for secure software supply chains
Docker in Edge Computing
Learn how to implement Docker in edge computing environments for efficient containerized applications at the network edge