From 571678492baee1fc22dbc19e39b3bcf9f3d28fd3 Mon Sep 17 00:00:00 2001 From: "Paul S. Schweigert" Date: Mon, 20 Jul 2026 22:50:37 -0400 Subject: [PATCH] dockerfile for running vllm with switch Signed-off-by: Paul S. Schweigert --- .dockerignore | 52 +++++++++++++++++++++++++++++++++++++++++++ Dockerfile | 61 +++++++++++++++++++++++++++++++++++++++++++++++++++ Makefile | 25 ++++++++++++++++++++- README.md | 21 ++++++++++++++++++ 4 files changed, 158 insertions(+), 1 deletion(-) create mode 100644 .dockerignore create mode 100644 Dockerfile diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..f83d591 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,52 @@ +# Only src/, pyproject.toml, and README.md are needed to build the wheel. +# Keep the build context small so `docker build` is fast and cache-stable. + +# VCS +.git/ +.github/ + +# Everything not needed to build/install the package +tests/ +tutorials/ +docs/ +scratch/ +scripts/ +tmp/ +.cache/ + +# Build artifacts +*.egg-info/ +.eggs/ +dist/ +build/ +*.egg + +# Python caches +__pycache__/ +*.py[cod] +*$py.class + +# Virtual envs +venv/ +.venv/ + +# Model/adapter checkpoints (multi-GB — must never enter the build context) +modular-granite/ +my-modular-model/ + +# Local tooling / editor / OS +.vscode/ +.idea/ +.DS_Store +.python-version +.claude/ + +# Docker's own files +Dockerfile +.dockerignore + +# Logs / coverage +*.log +.coverage +.coverage.* +htmlcov/ diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..a25c998 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: Apache-2.0 +# +# Serve Granite Switch with vLLM's OpenAI-compatible API server. +# +# Built on the official vLLM image, which already ships vLLM 0.19.1, PyTorch, +# and CUDA 12.x — matching this repo's default (vLLM 0.20+ needs CUDA 13+). +# We only add the granite_switch package on top; vLLM auto-discovers the model +# via the `vllm.general_plugins` entry point (see pyproject.toml), so no manual +# registration is needed at runtime. +# +# Build: +# docker build -t granite-switch-vllm . +# +# Run (needs an NVIDIA GPU + the NVIDIA container toolkit): +# docker run --gpus all -p 8000:8000 \ +# -v ~/.cache/huggingface:/root/.cache/huggingface \ +# granite-switch-vllm +# +# Serve a different model / checkpoint: +# docker run --gpus all -p 8000:8000 \ +# -e MODEL=ibm-granite/granite-switch-4.1-8b-preview \ +# granite-switch-vllm +# +# Pass extra vLLM args (appended after the entrypoint): +# docker run --gpus all -p 8000:8000 granite-switch-vllm \ +# --tensor-parallel-size 2 --max-model-len 8192 +# +# Gated/private models: pass a token with `-e HF_TOKEN=...`. + +# Pin to CUDA 12.x / vLLM 0.19.1. The plain tag is the x86_64 CUDA 12 build; +# the `-cu130` variants target CUDA 13. Bump in lockstep with pyproject's +# vllm pin. +FROM vllm/vllm-openai:v0.19.1 + +# Serve the 3b preview by default; override with `-e MODEL=...`. +ENV MODEL=ibm-granite/granite-switch-4.1-3b-preview \ + HOST=0.0.0.0 \ + PORT=8000 \ + # Faster HF downloads on first run (base image includes hf_transfer). + HF_HUB_ENABLE_HF_TRANSFER=1 + +WORKDIR /app + +# Install granite_switch itself. vLLM/torch/CUDA already live in the base +# image, so we deliberately install WITHOUT the [vllm] extra to avoid pulling +# a second, possibly conflicting vLLM wheel. Copy only what's needed to build +# the wheel so edits to tests/docs don't bust this layer's cache. +COPY pyproject.toml README.md ./ +COPY src/ ./src/ +RUN --mount=type=cache,target=/root/.cache/pip \ + pip install --no-deps . \ + && python -c "import granite_switch.vllm # sanity: package imports" + +EXPOSE 8000 + +# The base image's entrypoint is the vLLM API server. We can't reference $MODEL +# from an exec-form ENTRYPOINT, so wrap the launch in a tiny shell that expands +# the env vars and forwards any extra CLI args ("$@") to vLLM. +ENTRYPOINT ["/bin/bash", "-c", \ + "exec python3 -m vllm.entrypoints.openai.api_server --model \"$MODEL\" --host \"$HOST\" --port \"$PORT\" \"$@\"", \ + "--"] diff --git a/Makefile b/Makefile index 5eb2633..d1931e3 100644 --- a/Makefile +++ b/Makefile @@ -1,4 +1,4 @@ -.PHONY: test test-unit test-composer test-hf test-vllm test-integration test-all test-cpu test-gpu test-gpu-full test-tp test-regression-fast test-regression-real test-regression-hf test-peft-equiv test-peft-equiv-hf lint help +.PHONY: test test-unit test-composer test-hf test-vllm test-integration test-all test-cpu test-gpu test-gpu-full test-tp test-regression-fast test-regression-real test-regression-hf test-peft-equiv test-peft-equiv-hf lint docker-build docker-serve help # Default Python (can override: make test PYTHON=python3.11) PYTHON ?= python @@ -6,6 +6,13 @@ PYTHON ?= python # Pytest flags per CLAUDE.md guidelines PYTEST_FLAGS = -v -s --tb=short +# Docker image + serving config (override on the command line, e.g. +# `make docker-serve MODEL=ibm-granite/granite-switch-4.1-8b-preview PORT=8001`) +IMAGE ?= granite-switch-vllm +MODEL ?= ibm-granite/granite-switch-4.1-3b-preview +PORT ?= 8000 +HF_CACHE ?= $(HOME)/.cache/huggingface + # Individual test suites test-unit: $(PYTHON) -m pytest tests/unit/ $(PYTEST_FLAGS) @@ -63,6 +70,20 @@ test-peft-equiv-hf: lint: $(PYTHON) -m ruff check . +# Docker: build the vLLM serving image +docker-build: + docker build -t $(IMAGE) . + +# Docker: serve MODEL on PORT (requires an NVIDIA GPU + container toolkit). +# Mounts the host HF cache so models download once. Pass HF_TOKEN=... for +# gated models; extra vLLM args go in ARGS (e.g. ARGS="--tensor-parallel-size 2"). +docker-serve: + docker run --rm --gpus all -p $(PORT):$(PORT) \ + -v $(HF_CACHE):/root/.cache/huggingface \ + $(if $(HF_TOKEN),-e HF_TOKEN=$(HF_TOKEN),) \ + -e MODEL=$(MODEL) -e PORT=$(PORT) \ + $(IMAGE) $(ARGS) + # Help help: @echo "Available targets:" @@ -83,3 +104,5 @@ help: @echo " test-regression-hf - Run all HF regression tests (no vLLM/GPU)" @echo " test-peft-equiv-hf - Run PEFT equivalence (HF only)" @echo " lint - Run ruff linter" + @echo " docker-build - Build the vLLM serving image (override IMAGE=...)" + @echo " docker-serve - Serve MODEL on PORT via vLLM (override MODEL/PORT/ARGS/HF_TOKEN)" diff --git a/README.md b/README.md index 15fa6f8..466792f 100644 --- a/README.md +++ b/README.md @@ -115,6 +115,27 @@ print(f"social_bias score: {score:.3f}") # => social_bias score: 0.964 ``` +**Deploy with Docker:** + +The included `Dockerfile` builds on the official vLLM image and serves the OpenAI-compatible +API server. It requires an NVIDIA GPU and the [NVIDIA container toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html). + +```bash +make docker-build # build the image +make docker-serve # serve granite-switch-4.1-3b-preview on :8000 +``` + +The host Hugging Face cache (`~/.cache/huggingface`) is mounted so models download only once. +Override the model, port, or vLLM args on the command line: + +```bash +make docker-serve MODEL=ibm-granite/granite-switch-4.1-8b-preview PORT=8001 \ + ARGS="--tensor-parallel-size 2" +``` + +For gated models, pass a token with `HF_TOKEN=...`. The server then works with the same +Mellea client shown above (point `base_url` at the host and port you exposed). + ## How It Works With standard LoRA, each adapter is trained against its own KV distribution — so switching adapter functions across complex flow control means discarding and recomputing the KV cache at every step. aLoRA adapter functions are instead trained against a common normalized KV cache, so they can all coexist in a single checkpoint and activate on demand without cross-contamination: