54.11 GPU containers, Jupyter and a FastAPI model server
Friends, machine learning teams use Docker every day. A data scientist needs Jupyter with the right Python libraries, a training job needs the GPU, and the trained model must be served as an API. Today we do all three step by step: the NVIDIA Container Toolkit, --gpus, CUDA base images, a Jupyter container with a token, and a FastAPI model server that loads the model once.
मित्रांनो, machine learning teams रोज Docker वापरतात. Data scientist ला योग्य Python libraries सह Jupyter लागतो, training job ला GPU लागतो, आणि train केलेला model API म्हणून serve करावा लागतो. आज तिन्ही step by step करूया: NVIDIA Container Toolkit, --gpus, CUDA base images, token सह Jupyter container, आणि model एकदाच load करणारा FastAPI model server.
दोस्तों, machine learning teams रोज़ Docker इस्तेमाल करती हैं. Data scientist को सही Python libraries के साथ Jupyter चाहिए, training job को GPU चाहिए, और train किया हुआ model API के रूप में serve करना होता है. आज तीनों step by step करेंगे: NVIDIA Container Toolkit, --gpus, CUDA base images, token के साथ Jupyter container, और model एक बार load करने वाला FastAPI model server.
1. Start from a story
At Meesho, a data science team builds a model that predicts whether a product image is blurry. Each scientist set up CUDA, drivers and PyTorch on their own laptop or EC2 machine, and every setup was slightly different.
Training worked on one machine and crashed on another with CUDA version errors. A notebook server was started with no token and port 8888 open to the internet.
The first model API loaded the model on every request, so each call took two seconds.
Neha, the ML platform engineer, standardised it: the NVIDIA Container Toolkit on GPU hosts, one pinned PyTorch CUDA image, Jupyter bound to localhost with a token, and a FastAPI server that loads the model once at startup. Today we follow her steps.
2. GPUs in containers
Step 1: Install the NVIDIA Container Toolkit
Why: Once per GPU host.
What: The driver lives on the host. The toolkit lets Docker pass the GPU into containers. On AWS, g4dn and g5 instances have NVIDIA GPUs.
Example 1: Check the driver
Type this:
nvidia-smi --query-gpu=name,driver_version --format=csv,noheader
You will see:
Tesla T4, 550.90.07
What it means: A g4dn.xlarge on AWS has one T4.
Example 2: Install toolkit
Type this:
sudo dnf install -y nvidia-container-toolkit >/dev/null && sudo nvidia-ctk runtime configure --runtime=docker && sudo systemctl restart docker && echo configured
You will see:
INFO[0000] Wrote updated config to /etc/docker/daemon.json
configured
What it means: On Ubuntu use apt with NVIDIA's repository.
Example 3: Runtime registered
Type this:
docker info --format '{{json .Runtimes}}' | grep -o nvidia
You will see:
nvidia
What it means: Docker now knows the nvidia runtime.
Step 2: Run a container with --gpus
Why: Give the GPU to a container.
What: --gpus all passes every GPU. You can pass one by index too.
Example 4: All GPUs
Type this:
docker run --rm --gpus all nvidia/cuda:12.4.1-base-ubuntu22.04 nvidia-smi -L
You will see:
GPU 0: Tesla T4 (UUID: GPU-6a1f...)
What it means: The container sees the GPU.
Example 5: One GPU
Type this:
docker run --rm --gpus '"device=0"' nvidia/cuda:12.4.1-base-ubuntu22.04 nvidia-smi --query-gpu=index --format=csv,noheader
You will see:
0
What it means: Useful on 4 or 8 GPU machines.
Example 6: Without the flag
Type this:
docker run --rm nvidia/cuda:12.4.1-base-ubuntu22.04 nvidia-smi 2>&1 | head -1
You will see:
exec: "nvidia-smi": executable file not found in $PATH
What it means: No --gpus, no GPU.
Step 3: CUDA image types
Why: base, runtime, devel.
What: base has the minimum, runtime adds CUDA libraries, devel adds compilers for building extensions. Use the smallest that works.
Example 7: Sizes
Type this:
for t in base runtime devel; do docker pull -q nvidia/cuda:12.4.1-$t-ubuntu22.04 >/dev/null; done; docker images nvidia/cuda --format '{{.Tag}} {{.Size}}'
You will see:
12.4.1-base-ubuntu22.04 243MB
12.4.1-runtime-ubuntu22.04 2.6GB
12.4.1-devel-ubuntu22.04 7.1GB
What it means: devel is huge; use it only in a build stage.
Example 8: Driver rule
Type this:
printf 'The host driver must be new enough for the CUDA version in the image. Driver 550 supports CUDA 12.4.\n'
You will see:
The host driver must be new enough for the CUDA version in the image. Driver 550 supports CUDA 12.4.
What it means: The root cause of Meesho's crashes.
Example 9: Check inside
Type this:
docker run --rm --gpus all nvidia/cuda:12.4.1-base-ubuntu22.04 nvidia-smi | grep -o 'CUDA Version: [0-9.]*'
You will see:
CUDA Version: 12.4
What it means: Driver and image agree.
Step 4: PyTorch on the GPU
Why: One pinned image for the whole team.
What: The official pytorch image bundles CUDA and cuDNN. Pin the full tag so everyone gets the same versions.
Example 10: Check CUDA from Python
Type this:
docker run --rm --gpus all pytorch/pytorch:2.4.1-cuda12.4-cudnn9-runtime python -c 'import torch;print(torch.__version__, torch.cuda.is_available(), torch.cuda.get_device_name(0))'
You will see:
2.4.1 True Tesla T4
What it means: PyTorch sees the GPU.
Example 11: Quick speed test
Type this:
docker run --rm --gpus all pytorch/pytorch:2.4.1-cuda12.4-cudnn9-runtime python -c 'import torch,time;a=torch.rand(4000,4000,device="cuda");t=time.time();[a@a for _ in range(50)];torch.cuda.synchronize();print(round(time.time()-t,2),"s")'
You will see:
0.71 s
What it means: On CPU the same work takes about 40 seconds.
Example 12: Shared memory
Type this:
docker run --rm --gpus all --shm-size=2g pytorch/pytorch:2.4.1-cuda12.4-cudnn9-runtime df -h /dev/shm | tail -1
You will see:
shm 2.0G 0 2.0G 0% /dev/shm
What it means: DataLoader workers need more than the default 64 MB.
Step 5: GPUs in Compose
Why: Same request, in YAML.
What: Compose reserves GPUs under deploy.resources.reservations.devices.
Example 13: Service
Type this:
printf '%s\n' 'services:' ' train:' ' image: pytorch/pytorch:2.4.1-cuda12.4-cudnn9-runtime' ' command: python -c "import torch;print(torch.cuda.device_count())"' ' shm_size: 2gb' ' deploy:' ' resources:' ' reservations:' ' devices:' ' - driver: nvidia' ' count: all' ' capabilities: [gpu]' > compose.yaml && docker compose config -q && echo valid
You will see:
valid
What it means: capabilities: [gpu] is required.
Example 14: Run
Type this:
docker compose run --rm train
You will see:
1
What it means: One GPU visible.
Example 15: Monitor
Type this:
watch -n 2 nvidia-smi --query-gpu=utilization.gpu,memory.used --format=csv,noheader
You will see:
87 %, 6120 MiB
What it means: Watch usage from the host during training.
Step 6: A training job container
Why: Repeatable runs.
What: Mount the dataset read-only and an output folder for the model. The image holds code and libraries, not data.
Example 16: Train script
Type this:
printf '%s\n' 'import torch, torch.nn as nn' 'm = nn.Sequential(nn.Linear(10, 32), nn.ReLU(), nn.Linear(32, 2)).cuda()' 'opt = torch.optim.Adam(m.parameters())' 'for i in range(200):' ' x = torch.rand(64, 10).cuda(); y = (x.sum(1) > 5).long()' ' loss = nn.functional.cross_entropy(m(x), y); opt.zero_grad(); loss.backward(); opt.step()' 'torch.save(m.cpu().state_dict(), "/out/model.pt"); print("loss", round(loss.item(), 3))' > train.py && echo ok
You will see:
ok
What it means: A tiny model so the lab runs fast.
Example 17: Run it
Type this:
mkdir -p out && docker run --rm --gpus all -v "$PWD/train.py":/app/train.py:ro -v "$PWD/out":/out pytorch/pytorch:2.4.1-cuda12.4-cudnn9-runtime python /app/train.py
You will see:
loss 0.214
What it means: Code read-only, output to a mounted folder.
Example 18: Model saved
Type this:
ls -lh out
You will see:
-rw-r--r-- 1 root root 3.9K model.pt
What it means: Use --user $(id -u):$(id -g) so files belong to you, not root.
3. Jupyter and a model API
Step 7: A safe Jupyter container
Why: Token, localhost, your files.
What: The Jupyter Docker Stacks images run as user jovyan. Publish on 127.0.0.1 only and keep the token.
Example 19: Run
Type this:
docker run -d --name nb -p 127.0.0.1:8888:8888 -v "$PWD":/home/jovyan/work quay.io/jupyter/scipy-notebook:2024-10-07 && sleep 5 && docker logs nb 2>&1 | grep -m1 'token='
You will see:
http://127.0.0.1:8888/lab?token=3f9c1e...
What it means: Open the link in your browser. Never share the token.
Example 20: Who runs it
Type this:
docker exec nb id
You will see:
uid=1000(jovyan) gid=100(users) groups=100(users)
What it means: Not root.
Example 21: Remote server
Type this:
ssh -N -L 8888:127.0.0.1:8888 ec2-user@gpu-box.example.com
You will see:
(tunnel open, no output)
What it means: Reach a server notebook through SSH, not an open port.
Step 8: Jupyter with the GPU
Why: PyTorch notebooks on CUDA.
What: Use the pytorch-notebook image with a CUDA tag and add --gpus all.
Example 22: Run
Type this:
docker rm -f nb && docker run -d --name nb --gpus all -p 127.0.0.1:8888:8888 -v "$PWD":/home/jovyan/work quay.io/jupyter/pytorch-notebook:cuda12-2024-10-07 && echo started
You will see:
started
What it means: The cuda12 tag includes GPU builds of PyTorch.
Example 23: Check in the notebook
Type this:
docker exec nb python -c 'import torch;print(torch.cuda.is_available())'
You will see:
True
What it means: Notebook cells now run on the GPU.
Example 24: Extra packages
Type this:
printf '%s\n' 'FROM quay.io/jupyter/pytorch-notebook:cuda12-2024-10-07' 'RUN pip install --no-cache-dir timm==1.0.9 pillow==10.4.0' > Dockerfile.nb && docker build -q -f Dockerfile.nb -t team-nb:1.0 . && echo built
You will see:
sha256:8d0e...
built
What it means: Bake team libraries into one image, not pip install in every notebook.
Step 9: A FastAPI model server
Why: Load once, predict many.
What: The lifespan function loads the model when the app starts. Each request only runs prediction.
Example 25: serve.py
Type this:
printf '%s\n' 'import torch, torch.nn as nn' 'from contextlib import asynccontextmanager' 'from fastapi import FastAPI' 'from pydantic import BaseModel' 'M = {}' '@asynccontextmanager' 'async def lifespan(app):' ' m = nn.Sequential(nn.Linear(10, 32), nn.ReLU(), nn.Linear(32, 2))' ' m.load_state_dict(torch.load("/models/model.pt", map_location="cpu")); m.eval(); M["m"] = m' ' yield' 'app = FastAPI(lifespan=lifespan)' 'class In(BaseModel):' ' x: list[float]' '@app.get("/health")' 'def health(): return {"ok": "m" in M}' '@app.post("/predict")' 'def predict(i: In):' ' with torch.no_grad(): return {"class": int(M["m"](torch.tensor([i.x])).argmax())}' > serve.py && python3 -m py_compile serve.py && echo ok
You will see:
ok
What it means: The model is loaded once in lifespan.
Example 26: Dockerfile
Type this:
printf '%s\n' 'FROM python:3.12-slim' 'RUN pip install --no-cache-dir torch==2.4.1 --index-url https://download.pytorch.org/whl/cpu && pip install --no-cache-dir fastapi==0.115.0 uvicorn==0.30.6' 'RUN useradd -r -u 10001 app' 'WORKDIR /app' 'COPY serve.py .' 'COPY out/model.pt /models/model.pt' 'USER 10001' 'EXPOSE 8000' 'CMD ["uvicorn", "serve:app", "--host", "0.0.0.0", "--port", "8000"]' > Dockerfile && docker build -q -t blur-api:1.0 . && docker images blur-api:1.0 --format '{{.Size}}'
You will see:
912MB
What it means: CPU-only torch keeps it under 1 GB; the CUDA build is about 5 GB.
Example 27: Predict
Type this:
docker run -d -p 8000:8000 --name api blur-api:1.0 && sleep 4 && curl -s -X POST localhost:8000/predict -H 'Content-Type: application/json' -d '{"x":[0.9,0.8,0.7,0.9,0.6,0.8,0.9,0.7,0.8,0.9]}'
You will see:
{"class":1}
What it means: Fast, because the model is already in memory.
Step 10: Health, workers and timing
Why: Make the API production-like.
What: Add a healthcheck, choose workers carefully because each worker loads its own copy of the model, and measure latency.
Example 28: Healthcheck
Type this:
printf '%s\n' 'HEALTHCHECK --interval=15s --start-period=20s CMD python -c "import urllib.request;urllib.request.urlopen(\x27http://127.0.0.1:8000/health\x27)"' >> Dockerfile && docker build -q -t blur-api:1.1 . && echo built
You will see:
sha256:2a77...
built
What it means: start-period gives time to load the model.
Example 29: Latency
Type this:
for i in 1 2 3; do curl -s -o /dev/null -w '%{time_total}\n' -X POST localhost:8000/predict -H 'Content-Type: application/json' -d '{"x":[0,0,0,0,0,0,0,0,0,1]}'; done
You will see:
0.004
0.003
0.003
What it means: 3 to 4 ms, not the 2 seconds from the story.
Example 30: Workers and memory
Type this:
docker run -d -p 8001:8000 --name api2 blur-api:1.1 uvicorn serve:app --host 0.0.0.0 --port 8000 --workers 2 && sleep 5 && docker stats --no-stream --format '{{.Name}} {{.MemUsage}}' api api2
You will see:
api 182MiB / 15.4GiB
api2 351MiB / 15.4GiB
What it means: Two workers, two model copies. Big models need fewer workers.
Step 11: Keep big models out of the image
Why: Volumes or downloads at start.
What: A 5 GB model inside the image makes every push and pull slow. Mount it from a volume, or download it from S3 at startup.
Example 31: Mount a model folder
Type this:
docker run -d -p 8002:8000 -v "$PWD/out":/models:ro --name api3 blur-api:1.1 && sleep 4 && curl -s localhost:8002/health
You will see:
{"ok":true}
What it means: Change the model without rebuilding.
Example 32: Download from S3 at start
Type this:
printf '%s\n' 'aws s3 cp s3://example-ml-models/blur/v3/model.pt /models/model.pt' 'exec uvicorn serve:app --host 0.0.0.0 --port 8000'
You will see:
aws s3 cp s3://example-ml-models/blur/v3/model.pt /models/model.pt
exec uvicorn serve:app --host 0.0.0.0 --port 8000
What it means: Put this in an entrypoint script; use an IAM role, not access keys.
Example 33: Hugging Face cache volume
Type this:
docker run --rm -v hf-cache:/root/.cache/huggingface python:3.12-slim sh -c 'ls /root/.cache/huggingface | wc -l'
You will see:
0
What it means: A named volume keeps downloaded models between runs.
Step 12: Clean up GPU and notebook containers
Why: Free memory and disk.
What: ML images are large. Stop containers, check GPU memory is free, and prune old images.
Example 34: Stop
Type this:
docker rm -f nb api api2 api3 2>/dev/null | wc -l
You will see:
4
What it means: All lab containers removed.
Example 35: GPU free
Type this:
nvidia-smi --query-gpu=memory.used --format=csv,noheader
You will see:
0 MiB
What it means: No container holds GPU memory now.
Example 36: Disk
Type this:
docker system df --format '{{.Type}} {{.Size}}' | head -1 && docker image prune -f | tail -1
You will see:
Images 24.6GB
Total reclaimed space: 7.1GB
What it means: Remove dangling layers regularly.
4. Common mistakes and how Neha avoids them
Mistake one: installing CUDA and drivers inside every image. Neha keeps the driver on the host and CUDA libraries in pinned images.
Mistake two: a CUDA version newer than the host driver supports. She checks nvidia-smi first.
Mistake three: Jupyter on 0.0.0.0 without a token. She binds to 127.0.0.1 and uses SSH tunnels.
Mistake four: loading the model on every request. She loads it once in lifespan.
Mistake five: putting huge models in the image. She mounts them or downloads them at start.
Check yourself
Q1. What gives Docker access to the GPU?
The NVIDIA Container Toolkit plus --gpus.
Q2. Where does the NVIDIA driver live?
On the host, not in the image.
Q3. Why bind Jupyter to 127.0.0.1?
So it is not reachable from the network; use an SSH tunnel.
Q4. Why load the model in lifespan?
So it loads once at startup, not on every request.
Friends, driver on the host, toolkit for Docker, --gpus to request GPUs, pinned CUDA and PyTorch images, Jupyter on localhost with a token, and a FastAPI server that loads the model once. Did you understand? Let us try it once more. Train the tiny model and serve it with FastAPI on your own machine.
मित्रांनो, driver host वर, Docker साठी toolkit, GPU मागण्यासाठी --gpus, pinned CUDA आणि PyTorch images, token सह localhost वर Jupyter, आणि model एकदाच load करणारा FastAPI server. समजलं का? चला, पुन्हा एकदा करून पाहू. छोटा model train करा आणि तुमच्या machine वर FastAPI ने serve करा.
दोस्तों, driver host पर, Docker के लिए toolkit, GPU माँगने के लिए --gpus, pinned CUDA और PyTorch images, token के साथ localhost पर Jupyter, और model एक बार load करने वाला FastAPI server. समझ आया? चलिए, एक बार फिर करके देखते हैं। छोटा model train कीजिए और अपनी machine पर FastAPI से serve कीजिए.
Ravindra Bagale's Tip
GPU machines are expensive. Stop your g4dn or g5 instance when training is done, and do the API and notebook practice on a normal CPU machine.
GPU machines महाग असतात. Training संपल्यावर g4dn किंवा g5 instance stop करा, आणि API व notebook चा सराव साध्या CPU machine वर करा.
GPU machines महँगी होती हैं. Training खत्म होते ही g4dn या g5 instance stop कीजिए, और API व notebook की practice साधारण CPU machine पर कीजिए.
Lab: try at home
- Install the NVIDIA Container Toolkit on a GPU host.
- Run nvidia-smi inside a CUDA base container.
- Check torch.cuda.is_available in the PyTorch image.
- Start Jupyter on 127.0.0.1 with a token.
- Train the tiny model with mounted folders.
- Serve it with FastAPI and measure latency.
Recap
- Driver on the host, toolkit for Docker.
- --gpus all or a device index.
- Pick base, runtime or devel CUDA images wisely.
- Jupyter on localhost with a token.
- Load the model once in FastAPI lifespan.
- Keep large models in volumes or object storage.
Lab
- Pick two stacks and build both images.
- Compare their sizes with docker images.
- Run each as non-root and check with docker exec id (or USER in inspect).
- Bring MySQL, Postgres and Redis up with the compose file and confirm all are healthy.
Practice task
Containerise one real app in your own language using this chapter's pattern: multi-stage, non-root, healthcheck. Then add it to a compose.yaml with Postgres and Redis.
Ravindra Bagale's Tip
Every stack follows the same pattern: compile in the build stage, keep only what you need in the runtime. Remember this pattern and it works even when the language changes.
प्रत्येक stack चा pattern एकच: build stage मध्ये compile, runtime मध्ये फक्त गरजेचं. हा pattern लक्षात ठेवा, भाषा बदलली तरी चालेल.
हर stack का pattern एक ही: build stage में compile, runtime में सिर्फ़ ज़रूरी चीज़ें. यह pattern याद रखिए, भाषा बदल जाए तब भी काम करेगा.
In short
- Every stack follows build stage then slim runtime.
- Run as non-root and expose a high port.
- Use the language's production server (gunicorn, JarLauncher, php-fpm).
- Databases need volumes, password files and healthchecks.
- Init scripts run only on an empty volume.
Revision card
Interview card