Build and push a NIM-compatible container

On day 53 I built a production Dockerfile for vLLM. Today I want to take the next step: packaging an inference engine into a container that conforms to NVIDIA's NIM (NVIDIA Inference Microservice) standard. NIM is not just a container image; it is a contract that defines how the container discovers its GPU, selects an optimized model profile, exposes health and inference endpoints, and integrates with NVIDIA's ecosystem.

What is the NIM contract?

A NIM container is expected to implement the following behaviors:

OpenAI-compatible API: Expose /v1/chat/completions, /v1/completions, and /v1/models endpoints that conform to the OpenAI API schema.
Health endpoints: /v1/health/ready returns 200 when the model is loaded and ready. /v1/health/live returns 200 as long as the process is running.
GPU auto-detection: On startup, the container detects the available GPU type and count, then selects the appropriate model profile (quantization level, tensor parallelism degree) automatically.
Model profiles: The container bundles multiple optimized versions of the model. For example, an FP16 profile for H100, an FP8 profile for H100 with less memory, and an INT8 profile for A100.
NGC authentication: Model weights are pulled from NVIDIA's NGC registry using an API key, not baked into the image.

The model profile system

This is the most interesting part of the NIM design. Instead of one model binary per container, a NIM contains a manifest that maps GPU types to optimized model profiles:

# Example NIM model manifest (model_manifest.yaml)
profiles:
  - name: "tensorrt-llm-fp16-h100"
    gpu: "H100"
    min_gpus: 1
    precision: "fp16"
    engine: "tensorrt-llm"
    max_batch_size: 128
    max_seq_len: 8192

  - name: "tensorrt-llm-fp8-h100"
    gpu: "H100"
    min_gpus: 1
    precision: "fp8"
    engine: "tensorrt-llm"
    max_batch_size: 256
    max_seq_len: 8192

  - name: "vllm-fp16-a100"
    gpu: "A100"
    min_gpus: 1
    precision: "fp16"
    engine: "vllm"
    max_batch_size: 64
    max_seq_len: 4096

  - name: "vllm-fp16-a100-tp2"
    gpu: "A100"
    min_gpus: 2
    precision: "fp16"
    engine: "vllm"
    tensor_parallel: 2
    max_batch_size: 128
    max_seq_len: 8192

When the container starts, it queries the GPU type via nvidia-smi or the CUDA runtime, counts the available GPUs, and selects the best matching profile. This means the same container image runs on an A100 with vLLM in FP16 or on an H100 with TensorRT-LLM in FP8, automatically.

Building a NIM-compatible container

Here is a simplified but functional NIM-compatible container structure. The entrypoint script handles GPU detection and profile selection:

#!/bin/bash
# nim-entrypoint.sh - GPU detection and profile selection

set -euo pipefail

# Detect GPU type
GPU_NAME=$(nvidia-smi --query-gpu=name --format=csv,noheader | head -1 | xargs)
GPU_COUNT=$(nvidia-smi --query-gpu=count --format=csv,noheader | head -1 | xargs)

echo "Detected GPU: ${GPU_NAME} x${GPU_COUNT}"

# Select profile based on GPU
if echo "$GPU_NAME" | grep -qi "H100"; then
    if [ "$GPU_COUNT" -ge 2 ]; then
        PROFILE="tensorrt-llm-fp8-h100-tp2"
    else
        PROFILE="tensorrt-llm-fp8-h100"
    fi
elif echo "$GPU_NAME" | grep -qi "A100"; then
    if [ "$GPU_COUNT" -ge 2 ]; then
        PROFILE="vllm-fp16-a100-tp2"
    else
        PROFILE="vllm-fp16-a100"
    fi
elif echo "$GPU_NAME" | grep -qi "L40"; then
    PROFILE="vllm-int8-l40s"
else
    echo "WARNING: Unknown GPU, falling back to vllm-fp16 generic"
    PROFILE="vllm-fp16-generic"
fi

# Allow manual override
PROFILE="${NIM_MODEL_PROFILE:-$PROFILE}"
echo "Selected profile: ${PROFILE}"

# Load profile config and start the engine
exec python3 /opt/nim/launch.py --profile "${PROFILE}"

The launch script

The launch script reads the profile configuration and starts the appropriate engine:

# /opt/nim/launch.py
import argparse
import subprocess
import yaml
import sys

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--profile", required=True)
    args = parser.parse_args()

    with open("/opt/nim/model_manifest.yaml") as f:
        manifest = yaml.safe_load(f)

    profile = None
    for p in manifest["profiles"]:
        if p["name"] == args.profile:
            profile = p
            break

    if profile is None:
        print(f"ERROR: Profile {args.profile} not found")
        sys.exit(1)

    engine = profile["engine"]

    if engine == "vllm":
        cmd = [
            "python3", "-m", "vllm.entrypoints.openai.api_server",
            "--model", profile.get("model_path", "/models/default"),
            "--host", "0.0.0.0",
            "--port", "8000",
            "--max-model-len", str(profile["max_seq_len"]),
            "--gpu-memory-utilization", "0.9",
        ]
        if profile.get("tensor_parallel", 1) > 1:
            cmd += ["--tensor-parallel-size", str(profile["tensor_parallel"])]
    elif engine == "tensorrt-llm":
        cmd = [
            "python3", "/opt/nim/trtllm_server.py",
            "--engine-dir", profile.get("engine_path", "/engines/default"),
            "--host", "0.0.0.0",
            "--port", "8000",
        ]

    print(f"Starting: {' '.join(cmd)}")
    subprocess.execvp(cmd[0], cmd)

if __name__ == "__main__":
    main()

The Dockerfile

FROM nvcr.io/nvidia/cuda:12.4.1-runtime-ubuntu22.04

ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONDONTWRITEBYTECODE=1
ENV PYTHONUNBUFFERED=1

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

RUN python3.11 -m venv /opt/vllm-env
ENV PATH="/opt/vllm-env/bin:$PATH"

RUN pip install --no-cache-dir vllm==0.6.4 pyyaml

# NIM structure
COPY nim-entrypoint.sh /opt/nim/entrypoint.sh
COPY launch.py /opt/nim/launch.py
COPY model_manifest.yaml /opt/nim/model_manifest.yaml
RUN chmod +x /opt/nim/entrypoint.sh

# Health check using NIM standard endpoints
HEALTHCHECK --interval=30s --timeout=10s --start-period=180s --retries=3 \
    CMD curl -f http://localhost:8000/v1/health/ready || exit 1

EXPOSE 8000

ENTRYPOINT ["/opt/nim/entrypoint.sh"]

Pushing to a registry

# Build the NIM container
docker build -t my-nim:latest .

# Tag for your registry
docker tag my-nim:latest myregistry.azurecr.io/nim/llama-8b:v1.0

# Push
docker push myregistry.azurecr.io/nim/llama-8b:v1.0

# Run on any GPU - it auto-selects the right profile
docker run --gpus all -p 8000:8000 myregistry.azurecr.io/nim/llama-8b:v1.0

Why NIM matters for enterprises

The NIM contract separates "who optimizes the model" from "who deploys it." A platform team can deploy NIM containers without knowing the details of TensorRT-LLM compilation or vLLM configuration. The container handles GPU detection, profile selection, and engine configuration automatically. This is the same separation of concerns that made Docker successful for web applications.

A NIM container is a promise: give me a GPU and I will figure out the rest. That abstraction is what makes GPU inference deployable by platform teams, not just ML engineers.

After a day off, I will look at measuring cold start latency, the metric that determines how fast your autoscaling actually responds to load spikes.