NVIDIA NIMs on the Cluster

Last updated November 09, 2023
Table of Contents

NVIDIA Inference Microservices (NIMs) are a set of pre‑packaged, production‑ready containers that expose state‑of‑the‑art AI models (LLMs, vision, speech, embeddings, etc.) as HTTP/REST or gRPC services. The containers bundle the model, all required libraries (TensorRT, cuDNN, CUDA), and a tiny inference server that requests routing, batching, and metrics. You can spin‑up an inference endpoint with a single command, no need to write your own TensorRT‑or‑PyTorch serving code.

1 Prerequisites

  1. GPU: Any GPU, so you must submit your job to gpu partition or your condo node with a GPU
  2. Apptainer: Load the module with module load apptainer before running the container
  3. NGC / NIM account: Free registration at https://ngc.nvidia.com. You need an API key to pull private NIM images. Once you’ve signed up/in, go to Create, Setup, and copy the 32‑char string API key.
  4. Storage space: Models can take a significant amount of space (100GB+). Check your quota with myquota command.

1. Set API key (once per session)

export NGC_API_KEY="YOUR_32_CHAR_KEY"

2. Pull a NIM container (example: BERT)

mkdir -p /project2/ttrojan_123/apptainer_images
apptainer pull \
    --env NGC_API_KEY=$NGC_API_KEY \
    /project2/ttrojan_123/apptainer_images/bert_latest.sif \
    docker://nvcr.io/nim/bert:latest

3. Minimal Slurm script (run_bert.sbatch)

cat > run_bert.sbatch <<'EOF'
#!/bin/bash
#SBATCH --job-name=bert_nim
#SBATCH --partition=gpu
#SBATCH --gres=gpu:1
#SBATCH --cpus-per-task=4
#SBATCH --mem=32G
#SBATCH --time=02:00:00
#SBATCH --account=ttrojan_123
#SBATCH --output=bert_%j.out
#SBATCH --error=bert_%j.err

module purge
module load apptainer
#save the compute node IP address to the output file
ip a


CONTAINER=/project2/ttrojan_123/apptainer_images/bert_latest.sif
PORT=8500

apptainer exec --nv $CONTAINER \
    python -m nim_server --model_dir /opt/nim/models/bert \
    --port $PORT --log_level INFO
EOF

4. Submit

sbatch run_bert.sbatch

5. Test

Once the job is RUNNING, check the IP address in the output file under the the interface and use it to communicate with the NIM.

NODE=$(squeue -j <jobid> -h -o "%R")
IP=$(getent hosts $NODE | awk '{print $1}')
curl -X POST http://${IP}:8500/v1/embeddings \
     -H "Content-Type: application/json" \
     -d '{"input":"Hello world"}'