dockerfile for running vllm with switch

.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/

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\" \"$@\"", \
+    "--"]

Makefile

@@ -1,11 +1,18 @@
-.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
 
 # 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)"

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: