diff --git a/docs/guides/vllm_serving.md b/docs/guides/vllm_serving.md new file mode 100644 index 00000000..336048ab --- /dev/null +++ b/docs/guides/vllm_serving.md @@ -0,0 +1,88 @@ +# Serving KerasHub models with vLLM using Kinetic + +## Overview + +Export a KerasHub model to the Hugging Face Transformers format and serve it +with [vLLM](https://docs.vllm.ai) — on a Cloud**TPU** or **GPU**, in a single +Kinetic job, with any Keras backend. KerasHub +causal LMs export natively with `export_to_transformers()`, producing a standard +Hugging Face checkpoint (config, safetensors weights, tokenizer) that is +independent of the backend used to create it. Any preset with a Transformers +exporter works (Gemma, Gemma 3, Qwen, GPT-2, …). + +The export runs in a short-lived child process: when it exits, the OS releases +its memory and the device is clean for vLLM. This keeps the export and serving +stacks isolated — so any Keras backend works — while staying a single Kinetic +job (the export is just a subprocess inside the pod). + +## Prerequisites + +1. **A node pool** for your accelerator (default scale-to-zero). Pick any + slice/GPU that fits your model: + + ```bash + kinetic pool add --accelerator tpu-v5litepod-1 --project your-project-id # TPU + kinetic pool add --accelerator gpu-l4 --project your-project-id # GPU + ``` + + Larger models need a bigger slice (e.g. `tpu-v5litepod-4`, `gpu-a100`). + Make sure your project has matching **quota**, or the pod will sit `Pending`. + +2. **A `requirements.txt`** next to the scripts. The base set depends on the + device; a non-default export backend adds one entry: + + | Device | base requirements | backend extras | + |--------|-----------------------------|------------------------------------------------------------------| + | TPU | `keras keras-hub vllm-tpu` | `jax` / `torch`: none · `tensorflow`: add `tensorflow` | + | GPU | `keras keras-hub vllm` | `torch`: none · `jax`: add `jax[cuda12]` · `tensorflow`: add `tensorflow[and-cuda]` | + + The default backend per device (`jax` on TPU, `torch` on GPU) needs no extras. + On TPU, `torch`/`tensorflow` exports run on the host CPU, leaving the chip for + vLLM. Kinetic builds the remote container to match your **local Python + version**, so use one with `vllm`/`vllm-tpu` wheels available (3.10–3.12). + +3. **Kaggle credentials** for gated models like Gemma: accept the + [license](https://www.kaggle.com/models/keras/gemma3) and set + `KAGGLE_USERNAME` / `KAGGLE_KEY` locally. Kinetic forwards them via + `capture_env_vars`. + +## The example + +Two files: `vllm_serving.py` (the orchestrator — set `DEVICE`, `BACKEND`, +`MODEL_PRESET`, and `ACCELERATOR` at the top) and `export_worker.py` (the +standalone export, run as a subprocess). + +```{literalinclude} ../../examples/export_worker.py +:language: python +:caption: examples/export_worker.py +``` + +```{literalinclude} ../../examples/vllm_serving.py +:language: python +:caption: examples/vllm_serving.py +``` + +On TPU the example sets `JAX_PLATFORMS=tpu,cpu` and runs the engine in-process +(`VLLM_ENABLE_V1_MULTIPROCESSING=0`); on GPU it exposes the NVIDIA driver before +the export. `find_spec("export_worker")` locates the worker on the pod (Kinetic +unpacks the job's files onto `sys.path`) without importing it. + +## Running + +```bash +python vllm_serving.py +``` + +The first run builds the container image (15–25 minutes; later runs reuse it as +a cache hit) and provisions a node from the scale-to-zero pool. Monitor from a +second terminal with `kinetic jobs list` and +`kinetic jobs logs --follow JOB_ID --project your-project-id`. + +## Single job vs. two jobs + +This example is one Kinetic job — one pod — that runs the export as a child +process (two processes, one pod). To **export once and serve the same checkpoint +many times**, split it into two jobs instead: an export job that uploads the +checkpoint to `KINETIC_OUTPUT_DIR` (a GCS path), and a serve job that downloads +and serves it. Each job gets its own container, which also lets `vllm` and +`vllm-tpu` live in separate images. diff --git a/docs/index.rst b/docs/index.rst index 4ad652a6..ccf0de14 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -23,6 +23,8 @@ Kinetic: Run ML workloads on cloud TPUs and GPUs guides/cost_optimization guides/distributed_training guides/containers + guides/vllm_tpu + guides/vllm_serving guides/advanced .. toctree:: diff --git a/examples/export_worker.py b/examples/export_worker.py new file mode 100644 index 00000000..13f6eae7 --- /dev/null +++ b/examples/export_worker.py @@ -0,0 +1,18 @@ +"""KerasHub -> Hugging Face export. Run as a subprocess by vllm_serving.py so its +device memory is released on exit. Configured via environment variables.""" + +import os + +import keras_hub # KERAS_BACKEND is set by the parent before this import + +preset = os.environ["MODEL_PRESET"] +export_path = os.environ["EXPORT_PATH"] +dtype = os.environ.get("DTYPE", "bfloat16") + +print( + f"[export] backend={os.environ.get('KERAS_BACKEND')} preset={preset}", + flush=True, +) +lm = keras_hub.models.CausalLM.from_preset(preset, dtype=dtype) +lm.export_to_transformers(export_path) +print(f"[export] done -> {export_path}", flush=True) diff --git a/examples/vllm_serving.py b/examples/vllm_serving.py new file mode 100644 index 00000000..5b4de388 --- /dev/null +++ b/examples/vllm_serving.py @@ -0,0 +1,99 @@ +"""Export a KerasHub model and serve it with vLLM on a Cloud TPU or GPU, in a +single Kinetic job. See docs/guides/vllm_serving.md.""" + +import importlib.util +import os +import subprocess +import sys + +import kinetic + +DEVICE = "tpu" # "tpu" | "gpu" +BACKEND = ( + "jax" if DEVICE == "tpu" else "torch" +) # "jax" | "torch" | "tensorflow" +MODEL_PRESET = "gemma3_4b" +DTYPE = "bfloat16" +EXPORT_DIR = "/tmp/hf_export" + +# Any slice/GPU that fits your model works, e.g. "tpu-v5litepod-4", "gpu-a100". +ACCELERATOR = "tpu-v5litepod-1" if DEVICE == "tpu" else "gpu-l4" + + +def _setup_gpu_runtime(): + # Expose the GKE NVIDIA driver so the export child (which inherits + # LD_LIBRARY_PATH) and vLLM can use the GPU. + import ctypes + + nvidia_dirs = [ + d + for d in ("/usr/local/nvidia/lib64", "/usr/local/nvidia/lib") + if os.path.isdir(d) + ] + if nvidia_dirs: + prev = os.environ.get("LD_LIBRARY_PATH", "") + os.environ["LD_LIBRARY_PATH"] = ":".join( + nvidia_dirs + ([prev] if prev else []) + ) + for lib in ("libcuda.so.1", "libnvidia-ml.so.1"): + for root in nvidia_dirs: + if os.path.exists(os.path.join(root, lib)): + ctypes.CDLL(os.path.join(root, lib), mode=ctypes.RTLD_GLOBAL) + break + os.environ["VLLM_USE_FLASHINFER_SAMPLER"] = "0" + + +@kinetic.run( + accelerator=ACCELERATOR, + capture_env_vars=["KAGGLE_*", "GOOGLE_CLOUD_*"], +) +def export_and_serve(prompts): + if DEVICE == "gpu": + _setup_gpu_runtime() + + # Export in a child process; on exit the OS frees its device memory. find_spec + # locates the worker on the pod without importing it (which would defeat the + # isolation). + worker = importlib.util.find_spec("export_worker").origin + child_env = { + **os.environ, + "KERAS_BACKEND": BACKEND, + "MODEL_PRESET": MODEL_PRESET, + "EXPORT_PATH": EXPORT_DIR, + "DTYPE": DTYPE, + } + subprocess.run([sys.executable, worker], env=child_env, check=True) + + if DEVICE == "tpu": + os.environ["VLLM_TARGET_DEVICE"] = "tpu" + os.environ["JAX_PLATFORMS"] = "tpu,cpu" + os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = ( + "0" # in-process engine, no fork + ) + serve_kwargs = dict(max_model_len=1024, tensor_parallel_size=1) + else: + serve_kwargs = dict(max_model_len=1024, gpu_memory_utilization=0.85) + + from vllm import LLM, SamplingParams + + llm = LLM( + model=EXPORT_DIR, load_format="safetensors", dtype=DTYPE, **serve_kwargs + ) + sampling = SamplingParams(temperature=0.6, top_p=0.9, max_tokens=128) + outputs = llm.generate(prompts, sampling) + return [ + {"prompt": o.prompt, "completion": o.outputs[0].text.strip()} + for o in outputs + ] + + +if __name__ == "__main__": + prompts = [ + "The future of artificial intelligence will involve", + "A short recipe for a perfect weekend:", + "In one sentence, the theory of relativity says", + ] + for r in export_and_serve(prompts): + print("=" * 60) + print("Prompt:", r["prompt"]) + print("Completion:", r["completion"])