Posted on
Artificial Intelligence

Artificial Intelligence Container Best Practices

Author
  • User
    linuxbash
    Posts by this author
    Posts by this author

Artificial Intelligence Container Best Practices (for Linux Bash Users)

When your model trains fine on your laptop but falls apart on a teammate’s server, containers are the difference between “it works on my machine” and real, repeatable, high‑performance AI. This guide distills field‑tested best practices for building, running, and shipping AI workloads in containers on Linux—so you get reliable CUDA, fast I/O, and secure, reproducible pipelines.

What you’ll take away:

  • Exactly how to install a container runtime (Docker or Podman) and the NVIDIA Container Toolkit on apt, dnf, and zypper systems

  • A battle‑tested Dockerfile pattern for AI

  • Run flags and filesystem tricks that speed up training/serving and reduce surprises

  • Security and observability tips that productionize your setup

0) Prerequisites: Install a container runtime (+ optional GPU support)

Choose one runtime. Podman (rootless by default) is a great secure default; Docker has the broadest AI examples and docs. You can install both.

Install Podman:

  • apt (Ubuntu/Debian):
sudo apt-get update
sudo apt-get install -y podman
  • dnf (Fedora/RHEL family):
sudo dnf -y install podman
  • zypper (openSUSE/SLE):
sudo zypper refresh
sudo zypper install -y podman

Install Docker Engine:

  • apt (Ubuntu/Debian):
sudo apt-get update
sudo apt-get install -y ca-certificates curl gnupg
sudo install -m 0755 -d /etc/apt/keyrings
curl -fsSL https://download.docker.com/linux/$(. /etc/os-release; echo "$ID")/gpg | sudo gpg --dearmor -o /etc/apt/keyrings/docker.gpg
echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/$(. /etc/os-release; echo "$ID") $(. /etc/os-release; echo "$VERSION_CODENAME") stable" | sudo tee /etc/apt/sources.list.d/docker.list > /dev/null
sudo apt-get update
sudo apt-get install -y docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin
sudo usermod -aG docker $USER
newgrp docker
  • dnf (Fedora):
sudo dnf -y install dnf-plugins-core
sudo dnf config-manager --add-repo https://download.docker.com/linux/fedora/docker-ce.repo
sudo dnf -y install docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin
sudo systemctl enable --now docker
sudo usermod -aG docker $USER
newgrp docker
  • zypper (openSUSE/SLE; use distro packages):
sudo zypper refresh
sudo zypper install -y docker
sudo systemctl enable --now docker
sudo usermod -aG docker $USER
newgrp docker

Optional: NVIDIA Container Toolkit (for GPU access in containers). Ensure you already have the proprietary NVIDIA driver installed on the host.

  • apt (Ubuntu/Debian):
distribution=$(. /etc/os-release; echo ${ID}${VERSION_ID})
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
curl -fsSL https://nvidia.github.io/libnvidia-container/$distribution/libnvidia-container.list | \
  sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
  sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
sudo apt-get update
sudo apt-get install -y nvidia-container-toolkit
sudo nvidia-ctk runtime configure --runtime=docker
sudo systemctl restart docker
  • dnf (Fedora/RHEL family):
distribution=$(. /etc/os-release; echo ${ID}${VERSION_ID})
sudo curl -fsSL -o /etc/yum.repos.d/nvidia-container-toolkit.repo \
  https://nvidia.github.io/libnvidia-container/$distribution/libnvidia-container.repo
sudo dnf -y clean expire-cache
sudo dnf -y install nvidia-container-toolkit
sudo nvidia-ctk runtime configure --runtime=docker
sudo systemctl restart docker
  • zypper (openSUSE/SLE):
distribution=$(. /etc/os-release; echo ${ID}${VERSION_ID})
sudo zypper addrepo https://nvidia.github.io/libnvidia-container/$distribution/libnvidia-container.repo
sudo zypper refresh
sudo zypper install -y nvidia-container-toolkit
sudo nvidia-ctk runtime configure --runtime=docker
sudo systemctl restart docker

Quick GPU smoke test (Docker):

docker run --rm --gpus all nvidia/cuda:12.4.1-base-ubuntu22.04 nvidia-smi

Tip: Podman supports NVIDIA GPUs via OCI hooks from the same toolkit; consult NVIDIA’s Podman notes if you prefer Podman for GPU jobs.


1) Start with the right base image and pin everything

Why it matters:

  • CUDA/CuDNN and compiler mismatches cause “works yesterday, broken today.”

  • Pinned images and tools make builds deterministic across dev/CI/prod.

Use vendor ML images or CUDA‑aware bases and pin exact tags:

  • pytorch/pytorch

  • tensorflow/tensorflow

  • nvidia/cuda

Example minimal, reproducible Dockerfile:

# syntax=docker/dockerfile:1.6
FROM pytorch/pytorch:2.3.1-cuda12.1-cudnn8-runtime

ENV PYTHONDONTWRITEBYTECODE=1 \
    PIP_NO_CACHE_DIR=1 \
    PIP_DISABLE_PIP_VERSION_CHECK=1

# Create non-root user early for better layer caching
RUN useradd -m -u 1000 app
USER app
WORKDIR /app

# Copy only lockfiles first to maximize cache hits
COPY --chown=app:app requirements.txt .

# BuildKit cache speeds up pip installs
RUN --mount=type=cache,target=/home/app/.cache/pip \
    pip install --no-deps --requirement requirements.txt

# Then copy source
COPY --chown=app:app . .

# Optional healthcheck (simple CUDA probe)
HEALTHCHECK --interval=30s --timeout=5s --retries=5 \
  CMD python -c "import torch; import sys; sys.exit(0 if torch.cuda.is_available() else 1)"

Build with Docker or Podman:

DOCKER_BUILDKIT=1 docker build -t my-ai-app:0.1.0 .
# or
podman build -t my-ai-app:0.1.0 .

Pin dependencies with lockfiles:

  • pip: generate requirements.txt with hashes (pip-compile, poetry export, uv pip compile)

  • conda/mamba: conda-lock to freeze explicit versions

  • Always record CUDA/CuDNN/driver versions alongside your image tag in CHANGELOGs


2) Optimize runtime flags for performance and stability

AI workloads are IO- and memory‑intense. A few flags prevent common footguns:

  • Bigger shared memory for PyTorch DataLoader, DALI, or CV pipelines:

  • Avoid interprocess comm issues with NCCL:

  • Keep memlock unlimited for better perf on some stacks:

Example (training with all GPUs, 8G shm, NCCL friendly, non-root):

docker run --rm --gpus all \
  --shm-size=8g --ipc=host \
  --ulimit memlock=-1 --ulimit stack=67108864 \
  -u 1000:1000 \
  -e NCCL_P2P_DISABLE=0 -e NCCL_IB_DISABLE=0 \
  -e OMP_NUM_THREADS=8 \
  -v $PWD:/app \
  -w /app \
  my-ai-app:0.1.0 python train.py --epochs 10

Notes:

  • ipc=host avoids small /dev/shm defaults; alternatively tune --shm-size only.

  • For multi-socket servers, pin CPU/GPU affinity using taskset/numactl (host) or CUDA_VISIBLE_DEVICES inside the container.

  • For inference, set thread counts (OMP, MKL) and batch sizes explicitly to prevent oversubscription.


3) Keep images lean and builds fast

Why it matters:

  • Smaller images pull faster, start faster, and reduce CVE surface area.

  • Faster builds mean more iteration.

Tactics:

  • Prefer “-runtime” over “-devel” bases for serving; use multi-stage builds.

  • Clean up package caches within the same RUN layer.

  • Cache pip/conda artifacts with BuildKit mounts (as shown above).

  • Group layers logically so dependency layers rarely change.

Example multi-stage for custom CUDA ops:

# syntax=docker/dockerfile:1.6
FROM nvidia/cuda:12.1.0-devel-ubuntu22.04 AS build
RUN apt-get update && apt-get install -y --no-install-recommends build-essential python3-dev python3-pip \
 && rm -rf /var/lib/apt/lists/*
WORKDIR /w
COPY requirements-build.txt .
RUN --mount=type=cache,target=/root/.cache/pip pip install -r requirements-build.txt
COPY . .
RUN python setup.py bdist_wheel

FROM nvidia/cuda:12.1.0-runtime-ubuntu22.04
RUN useradd -m -u 1000 app
USER app
WORKDIR /app
COPY --from=build /w/dist/*.whl .
RUN --mount=type=cache,target=/home/app/.cache/pip pip install *.whl
COPY . .
CMD ["python", "serve.py"]

4) Treat data and caches as volumes, not image content

Why it matters:

  • Shipping datasets inside images explodes size and breaks caching.

  • Managed caches speed up cold starts (especially HF/Transformers, pip, and model weights).

Mount datasets read-only and centralize caches:

mkdir -p /data/datasets /data/cache/hf /data/cache/pip

docker run --rm --gpus all \
  -v /data/datasets:/datasets:ro \
  -v /data/cache/hf:/cache/hf \
  -v /data/cache/pip:/home/app/.cache/pip \
  -e HF_HOME=/cache/hf \
  my-ai-app:0.1.0 python train.py --data /datasets/imagenet

If you serve with Hugging Face Transformers, this prevents re-downloading and encourages warm starts across restarts and replicas.


5) Run as non-root and drop privileges by default

Why it matters:

  • AI containers often touch network, mounts, and GPUs—limit blast radius.

Do this:

  • Create and use a non-root user in the image (see Dockerfile above).

  • Grant only needed capabilities and keep the root filesystem read-only in production.

  • Never use --privileged unless you absolutely must.

Example secure run for an inference server on port 8080:

docker run --rm --gpus all \
  --cap-drop=ALL --cap-add=NET_BIND_SERVICE \
  --read-only -v /tmp -v /data/cache/hf:/cache/hf \
  -e HF_HOME=/cache/hf \
  -p 8080:8080 \
  -u 1000:1000 \
  my-ai-app:0.1.0 python serve.py --host 0.0.0.0 --port 8080

Add a healthcheck to surface readiness/liveness to orchestrators:

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

Real‑world patterns

  • Training job (resumeable, dataset mount, all GPUs):
docker run --rm --gpus all \
  --shm-size=8g --ipc=host \
  -v /data/datasets:/datasets:ro \
  -v /data/ckpts:/ckpts \
  -e CHECKPOINT_DIR=/ckpts/run-42 \
  my-ai-app:0.1.0 bash -lc 'python train.py --data /datasets --out $CHECKPOINT_DIR --resume if-exists'
  • Inference server (low-latency, fixed threads, model cache):
docker run --rm --gpus all \
  -e OMP_NUM_THREADS=4 -e TF_NUM_INTRAOP_THREADS=4 -e TF_NUM_INTEROP_THREADS=2 \
  -v /data/cache/hf:/cache/hf -e HF_HOME=/cache/hf \
  -p 8080:8080 \
  my-ai-app:0.1.0 python serve.py --host 0.0.0.0 --port 8080
  • Rootless Podman (CPU-only dev loop):
podman run --rm -it \
  -v $PWD:/app -w /app \
  localhost/my-ai-app:0.1.0 bash

Troubleshooting quick hits

  • nvidia-smi works on host but not in container:

    • Ensure nvidia-container-toolkit installed and configured; use docker run --gpus all ...
  • PyTorch sees 0 GPUs:

    • Check CUDA version in image vs. host driver; newer driver supports older CUDA, not vice versa.
  • DataLoader hangs:

    • Increase --shm-size or use --ipc=host.
  • Slow cold starts:

    • Externalize caches (HF_HOME, pip cache) to host volumes.

Conclusion and next steps (CTA)

Containers are your reliability and performance multiplier for AI: pin the stack, right-size runtime flags, mount data/caches, and lock down privileges. Your next step:

1) Install your runtime and (optionally) GPU toolkit using the commands above. 2) Copy the sample Dockerfile, pin your dependencies, and build an image. 3) Run with the provided flags, mount your datasets read-only, and add a healthcheck. 4) Automate with a Makefile or CI to guarantee reproducibility.

When you’re ready, extend this checklist into your team’s container template so every new model starts from a proven, production-grade base.