From e968c55e354f0a0c067e853f9a619bcb1de18a3a Mon Sep 17 00:00:00 2001 From: Dhiraj BM Date: Fri, 12 Jun 2026 05:27:23 +0530 Subject: [PATCH 1/3] Add guide: serving KerasHub models with vLLM using Kinetic --- docs/guides/vllm_serving.md | 138 ++++++++++++++++++++++++++++ docs/index.rst | 2 + examples/vllm_serving.py | 173 ++++++++++++++++++++++++++++++++++++ 3 files changed, 313 insertions(+) create mode 100644 docs/guides/vllm_serving.md create mode 100644 examples/vllm_serving.py diff --git a/docs/guides/vllm_serving.md b/docs/guides/vllm_serving.md new file mode 100644 index 00000000..67259e08 --- /dev/null +++ b/docs/guides/vllm_serving.md @@ -0,0 +1,138 @@ +# Serving KerasHub Models with vLLM + +This guide explains how to export a KerasHub model to the Hugging Face +Transformers format and serve it with vLLM using the Kinetic framework — +on a Cloud GPU in a single job, or on TPU as a two-job workflow. + +## Overview + +KerasHub causal LMs can be exported natively with +`export_to_transformers()`, producing a standard Hugging Face checkpoint +(config, safetensors weights, tokenizer). Kinetic lets you run the export +and vLLM serving in a single GPU job: the model is downloaded and exported +on the remote worker, the GPU is handed over to vLLM, and batched +completions are returned to your local machine. The exported checkpoint is +also archived to `KINETIC_OUTPUT_DIR`, so it can be served again anywhere +vLLM runs. + +The example uses `gemma3_4b` (~8 GB in bfloat16), which fits a single +NVIDIA L4. `gpt2_large_en` is an ungated alternative if you don't have +Kaggle access to Gemma (use `dtype="float32"` with it). + +## Prerequisites + +1. **Kinetic Cluster**: You need a provisioned Kinetic cluster with GPU + nodes (e.g., `l4`), created with the latest `keras-kinetic`: + + ```bash + kinetic pool add --accelerator l4 --project your-project-id + ``` + + GPU quota is the most common first-time blocker: in *IAM & Admin → + Quotas*, both **GPUs (all regions)** and **NVIDIA L4 GPUs** (regional) + must be ≥ 1. +2. **Kaggle Credentials**: If you are using gated models like Gemma 3, + you need a Kaggle account with the + [model license accepted](https://www.kaggle.com/models/keras/gemma3), + and `KAGGLE_USERNAME` / `KAGGLE_KEY` set in your local environment. + +## Configuration + +To run vLLM successfully on GPU via Kinetic, you need to handle +dependencies and environment variables properly. + +### 1. Dependencies + +Create a `requirements.txt` file in the directory of your script +containing: + +```text +keras +keras-hub +tensorflow-text +vllm +``` + +Kinetic will detect this file and build a container with vLLM installed. +`tensorflow-text` is required — +KerasHub tokenizers preprocess with `tf.data` on every backend (CPU-side +only). Use a **Python 3.12** local venv; the remote container matches your +local interpreter, and `tensorflow-text` doesn't publish wheels for the +newest Python yet. + +### 2. Environment Variables + +The example sets the following environment variables on the remote worker +to ensure correct execution: + +- `KERAS_BACKEND="torch"`: vLLM is PyTorch-based, and the torch backend + is the only one that releases VRAM cleanly after the export, so vLLM's + KV cache gets the full GPU. +- `VLLM_USE_FLASHINFER_SAMPLER="0"`: vLLM's FlashInfer sampler + JIT-compiles CUDA kernels with `nvcc`, which pip-only containers don't + have; this selects the native torch sampler instead. +- `LD_LIBRARY_PATH=/usr/local/nvidia/lib64:...`: GKE mounts the host + NVIDIA driver at `/usr/local/nvidia`; this makes `libcuda` / + `libnvidia-ml` visible to vLLM's spawned engine processes (the example + also preloads them into the main process via `ctypes`). + +These are set inside the remote function by the example itself. Kaggle +credentials are forwarded from your local environment via +`capture_env_vars` in the `@kinetic.run` decorator. + +## Example + +```{literalinclude} ../../examples/vllm_serving.py +:language: python +``` + +## Running the Example + +```bash +python3 vllm_serving.py +``` + +The first run builds the container image (15–25 minutes; subsequent runs +reuse it as a cache hit) and provisions a GPU node from the scale-to-zero +pool (~10 minutes including the image pull). Monitor from a second +terminal with `kinetic jobs list` and +`kinetic jobs logs --follow JOB_ID --project your-project-id`. + +## Serving on TPU + +The export step is accelerator-agnostic, but vLLM serving itself always +needs a GPU or TPU. TPU serving uses a different vLLM build (`vllm-tpu`) +and different environment variables than the GPU example above, and +Kinetic builds one container per script directory — `vllm` and `vllm-tpu` +cannot share an image. So on TPU, export and serving run as **two scripts +in two directories, each with its own `requirements.txt`**. The checkpoint +moves between them through GCS; nothing is downloaded to your local +machine, and no re-initialization of Kinetic is needed. + +1. **Export script** (directory A, `requirements.txt`: `keras`, + `keras-hub`, `tensorflow-text`) — the export half of this example, on + any accelerator (`"cpu"` works). It archives the checkpoint to + `KINETIC_OUTPUT_DIR` (a GCS path), which it returns; pass that path to + the serving script. +2. **Serving script** (directory B, `requirements.txt`: `vllm-tpu`) — + configured per + [Running vLLM on TPU with Kinetic](../guides/vllm_tpu.md): set + `VLLM_TARGET_DEVICE="tpu"`, `VLLM_USE_V1="0"`, and + `JAX_PLATFORMS="tpu,cpu"` locally and forward them with + `capture_env_vars`. The job downloads the archive from GCS to the + pod's local disk, unpacks it, and points `LLM(model=...)` at that + directory instead of a Hub model ID: + +```python +@kinetic.run( + accelerator="tpu-v5litepod-8", + capture_env_vars=["VLLM_*", "JAX_*"], +) +def serve_on_tpu(checkpoint_gs_path, prompts): + # Download + unpack the exported checkpoint from GCS to /tmp/hf_export + # (google.cloud.storage download, tarfile.extractall), then: + from vllm import LLM, SamplingParams + + llm = LLM(model="/tmp/hf_export", max_model_len=1024) + return llm.generate(prompts, SamplingParams(max_tokens=128)) +``` 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/vllm_serving.py b/examples/vllm_serving.py new file mode 100644 index 00000000..68013668 --- /dev/null +++ b/examples/vllm_serving.py @@ -0,0 +1,173 @@ +import os + +import kinetic + +# Any preset with a Transformers exporter works (Gemma 3 / Gemma / +# Qwen / GPT-2). +MODEL_PRESET = "gemma3_4b" + +# Load and export in the model's native precision; "float32" for GPT-2 +# or pre-Ampere GPUs. +DTYPE = "bfloat16" + +EXPORT_DIR = "/tmp/hf_export" + +# vLLM routes checkpoints via the `architectures` key in config.json. +# The Gemma 3 exporter writes it natively; the others don't yet, so we +# patch it after export. TODO: drop once all exporters in +# `keras_hub/src/utils/transformers/export/` write this key. +HF_ARCHITECTURES = { + "gpt2": "GPT2LMHeadModel", + "qwen2": "Qwen2ForCausalLM", + "gemma": "GemmaForCausalLM", + "gemma3_text": "Gemma3ForCausalLM", +} + + +def _ensure_architectures(export_dir): + """Add the `architectures` key to config.json if missing.""" + import json + + config_path = os.path.join(export_dir, "config.json") + with open(config_path) as f: + config = json.load(f) + if "architectures" not in config: + arch = HF_ARCHITECTURES.get(config.get("model_type")) + if arch is None: + raise ValueError( + f"Unknown model_type {config.get('model_type')!r}; " + "add it to HF_ARCHITECTURES." + ) + config["architectures"] = [arch] + with open(config_path, "w") as f: + json.dump(config, f, indent=2) + + +def _setup_gpu_runtime(): + """Expose the GKE-mounted NVIDIA driver to torch/vLLM and select + the torch-native sampler (pip-only images have no nvcc).""" + import ctypes + import glob + + print("GPU devices:", glob.glob("/dev/nvidia*")) + + # For vLLM's spawned engine processes, which read this at startup. + 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 this process, whose linker ignores late LD_LIBRARY_PATH edits. + for lib in ("libcuda.so.1", "libnvidia-ml.so.1"): + for root in ("/usr/local/nvidia/lib64", "/usr/local/nvidia/lib"): + path = os.path.join(root, lib) + if os.path.exists(path): + ctypes.CDLL(path, mode=ctypes.RTLD_GLOBAL) + print("preloaded:", path) + break + else: + print("NOT FOUND:", lib) + + os.environ["VLLM_USE_FLASHINFER_SAMPLER"] = "0" + + +def _persist_export(export_dir): + """Archive the exported checkpoint to the job's durable output dir.""" + import shutil + + output_dir = os.environ.get("KINETIC_OUTPUT_DIR") + if not output_dir: + return None + # Plain tar: model weights are high-entropy, gzip only costs time. + print("Archiving exported checkpoint (this can take a few minutes)...") + archive = shutil.make_archive("/tmp/hf_export_archive", "tar", export_dir) + dest = f"{output_dir.rstrip('/')}/hf_export.tar" + if dest.startswith("gs://"): + from google.cloud import storage + + bucket_name, _, blob_name = dest[5:].partition("/") + blob = storage.Client().bucket(bucket_name).blob(blob_name) + blob.upload_from_filename(archive) + else: + os.makedirs(output_dir, exist_ok=True) + shutil.copy(archive, dest) + return dest + + +@kinetic.run( + accelerator="gpu-l4", + capture_env_vars=["KAGGLE_*", "GOOGLE_CLOUD_*"], +) +def export_and_serve(prompts): + _setup_gpu_runtime() + + # torch is the only backend that frees VRAM for vLLM in the same + # process; keras reads this at first (transitive) import. + os.environ["KERAS_BACKEND"] = "torch" + + import gc + + import keras_hub + import torch + + print( + "Kaggle creds in pod:", + bool(os.environ.get("KAGGLE_USERNAME")) + and bool(os.environ.get("KAGGLE_KEY")), + ) + + print(f"Loading {MODEL_PRESET} from KerasHub ({DTYPE})...") + lm = keras_hub.models.CausalLM.from_preset(MODEL_PRESET, dtype=DTYPE) + + print("Exporting to Hugging Face Transformers format...") + lm.export_to_transformers(EXPORT_DIR) + _ensure_architectures(EXPORT_DIR) + + # Release VRAM before handing the GPU to vLLM. + del lm + gc.collect() + torch.cuda.empty_cache() + + from vllm import LLM, SamplingParams + + print("Starting vLLM engine from the exported checkpoint...") + llm = LLM( + model=EXPORT_DIR, + dtype=DTYPE, + gpu_memory_utilization=0.85, + max_model_len=1024, + enforce_eager=False, + ) + sampling_params = SamplingParams(temperature=0.6, top_p=0.9, max_tokens=128) + + # All prompts are processed in one continuous-batching pass. + outputs = llm.generate(prompts, sampling_params) + results = [ + {"prompt": o.prompt, "completion": o.outputs[0].text.strip()} + for o in outputs + ] + + artifact = _persist_export(EXPORT_DIR) + if artifact: + print(f"Exported model archived at: {artifact}") + + return results + + +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", + "The most underrated skill in software engineering is", + ] + for result in export_and_serve(prompts): + print("=" * 60) + print(f"Prompt: {result['prompt']}") + print(f"Completion: {result['completion']}") From db20c5b5d6ff27450550053c168cdd32bdf15a7c Mon Sep 17 00:00:00 2001 From: Dhiraj BM Date: Fri, 12 Jun 2026 07:02:22 +0530 Subject: [PATCH 2/3] Address review: upload checkpoint directory before vLLM init --- docs/guides/vllm_serving.md | 16 +++++------ examples/vllm_serving.py | 53 ++++++++++++++++++++++++++----------- 2 files changed, 46 insertions(+), 23 deletions(-) diff --git a/docs/guides/vllm_serving.md b/docs/guides/vllm_serving.md index 67259e08..9b066467 100644 --- a/docs/guides/vllm_serving.md +++ b/docs/guides/vllm_serving.md @@ -1,4 +1,4 @@ -# Serving KerasHub Models with vLLM +# Serving KerasHub Models with vLLM using Kinetic This guide explains how to export a KerasHub model to the Hugging Face Transformers format and serve it with vLLM using the Kinetic framework — @@ -12,7 +12,7 @@ KerasHub causal LMs can be exported natively with and vLLM serving in a single GPU job: the model is downloaded and exported on the remote worker, the GPU is handed over to vLLM, and batched completions are returned to your local machine. The exported checkpoint is -also archived to `KINETIC_OUTPUT_DIR`, so it can be served again anywhere +also uploaded to `KINETIC_OUTPUT_DIR`, so it can be served again anywhere vLLM runs. The example uses `gemma3_4b` (~8 GB in bfloat16), which fits a single @@ -111,7 +111,7 @@ machine, and no re-initialization of Kinetic is needed. 1. **Export script** (directory A, `requirements.txt`: `keras`, `keras-hub`, `tensorflow-text`) — the export half of this example, on - any accelerator (`"cpu"` works). It archives the checkpoint to + any accelerator (`"cpu"` works). It uploads the checkpoint directory to `KINETIC_OUTPUT_DIR` (a GCS path), which it returns; pass that path to the serving script. 2. **Serving script** (directory B, `requirements.txt`: `vllm-tpu`) — @@ -119,9 +119,9 @@ machine, and no re-initialization of Kinetic is needed. [Running vLLM on TPU with Kinetic](../guides/vllm_tpu.md): set `VLLM_TARGET_DEVICE="tpu"`, `VLLM_USE_V1="0"`, and `JAX_PLATFORMS="tpu,cpu"` locally and forward them with - `capture_env_vars`. The job downloads the archive from GCS to the - pod's local disk, unpacks it, and points `LLM(model=...)` at that - directory instead of a Hub model ID: + `capture_env_vars`. The job downloads the checkpoint directory from + GCS to the pod's local disk and points `LLM(model=...)` at it instead + of a Hub model ID: ```python @kinetic.run( @@ -129,8 +129,8 @@ machine, and no re-initialization of Kinetic is needed. capture_env_vars=["VLLM_*", "JAX_*"], ) def serve_on_tpu(checkpoint_gs_path, prompts): - # Download + unpack the exported checkpoint from GCS to /tmp/hf_export - # (google.cloud.storage download, tarfile.extractall), then: + # Download the exported checkpoint directory from GCS to + # /tmp/hf_export (google.cloud.storage transfer_manager), then: from vllm import LLM, SamplingParams llm = LLM(model="/tmp/hf_export", max_model_len=1024) diff --git a/examples/vllm_serving.py b/examples/vllm_serving.py index 68013668..78cbbd30 100644 --- a/examples/vllm_serving.py +++ b/examples/vllm_serving.py @@ -1,3 +1,9 @@ +"""Export a KerasHub causal LM to the Hugging Face Transformers format +and serve it with vLLM on a Kinetic GPU, in a single job. + +See the accompanying guide for prerequisites and configuration. +""" + import os import kinetic @@ -78,25 +84,38 @@ def _setup_gpu_runtime(): def _persist_export(export_dir): - """Archive the exported checkpoint to the job's durable output dir.""" + """Upload the exported checkpoint directory to the job's durable + output dir.""" import shutil output_dir = os.environ.get("KINETIC_OUTPUT_DIR") if not output_dir: return None - # Plain tar: model weights are high-entropy, gzip only costs time. - print("Archiving exported checkpoint (this can take a few minutes)...") - archive = shutil.make_archive("/tmp/hf_export_archive", "tar", export_dir) - dest = f"{output_dir.rstrip('/')}/hf_export.tar" + dest = f"{output_dir.rstrip('/')}/hf_export" if dest.startswith("gs://"): + # Direct directory upload: no tar copy on ephemeral disk, and + # parallel workers make it fast. from google.cloud import storage - - bucket_name, _, blob_name = dest[5:].partition("/") - blob = storage.Client().bucket(bucket_name).blob(blob_name) - blob.upload_from_filename(archive) + from google.cloud.storage import transfer_manager + + bucket_name, _, blob_prefix = dest[5:].partition("/") + bucket = storage.Client().bucket(bucket_name) + files = [] + for root, _, filenames in os.walk(export_dir): + for filename in filenames: + rel_path = os.path.relpath(os.path.join(root, filename), export_dir) + files.append(rel_path) + print("Uploading exported checkpoint to GCS...") + transfer_manager.upload_many_from_filenames( + bucket, + files, + source_directory=export_dir, + blob_name_prefix=blob_prefix + "/", + worker_type=transfer_manager.THREAD, + raise_exception=True, + ) else: - os.makedirs(output_dir, exist_ok=True) - shutil.copy(archive, dest) + shutil.copytree(export_dir, dest, dirs_exist_ok=True) return dest @@ -116,6 +135,8 @@ def export_and_serve(prompts): import keras_hub import torch + # Forwarded via capture_env_vars; ~/.kaggle/kaggle.json doesn't + # travel to the pod. print( "Kaggle creds in pod:", bool(os.environ.get("KAGGLE_USERNAME")) @@ -129,6 +150,12 @@ def export_and_serve(prompts): lm.export_to_transformers(EXPORT_DIR) _ensure_architectures(EXPORT_DIR) + # Persist before starting vLLM: engine init is the riskiest step, + # and the checkpoint survives even if it crashes. + artifact = _persist_export(EXPORT_DIR) + if artifact: + print(f"Exported model uploaded to: {artifact}") + # Release VRAM before handing the GPU to vLLM. del lm gc.collect() @@ -153,10 +180,6 @@ def export_and_serve(prompts): for o in outputs ] - artifact = _persist_export(EXPORT_DIR) - if artifact: - print(f"Exported model archived at: {artifact}") - return results From d53f9f56836d0b9148589d44f2e0ec5436ae1529 Mon Sep 17 00:00:00 2001 From: Dhiraj BM Date: Wed, 17 Jun 2026 22:39:07 +0530 Subject: [PATCH 3/3] Updated the guide to support all backends --- docs/guides/vllm_serving.md | 172 +++++++++++-------------------- examples/export_worker.py | 18 ++++ examples/vllm_serving.py | 195 +++++++++--------------------------- 3 files changed, 128 insertions(+), 257 deletions(-) create mode 100644 examples/export_worker.py diff --git a/docs/guides/vllm_serving.md b/docs/guides/vllm_serving.md index 9b066467..336048ab 100644 --- a/docs/guides/vllm_serving.md +++ b/docs/guides/vllm_serving.md @@ -1,138 +1,88 @@ -# Serving KerasHub Models with vLLM using Kinetic - -This guide explains how to export a KerasHub model to the Hugging Face -Transformers format and serve it with vLLM using the Kinetic framework — -on a Cloud GPU in a single job, or on TPU as a two-job workflow. +# Serving KerasHub models with vLLM using Kinetic ## Overview -KerasHub causal LMs can be exported natively with -`export_to_transformers()`, producing a standard Hugging Face checkpoint -(config, safetensors weights, tokenizer). Kinetic lets you run the export -and vLLM serving in a single GPU job: the model is downloaded and exported -on the remote worker, the GPU is handed over to vLLM, and batched -completions are returned to your local machine. The exported checkpoint is -also uploaded to `KINETIC_OUTPUT_DIR`, so it can be served again anywhere -vLLM runs. +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 example uses `gemma3_4b` (~8 GB in bfloat16), which fits a single -NVIDIA L4. `gpt2_large_en` is an ungated alternative if you don't have -Kaggle access to Gemma (use `dtype="float32"` with it). +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. **Kinetic Cluster**: You need a provisioned Kinetic cluster with GPU - nodes (e.g., `l4`), created with the latest `keras-kinetic`: - - ```bash - kinetic pool add --accelerator l4 --project your-project-id - ``` - - GPU quota is the most common first-time blocker: in *IAM & Admin → - Quotas*, both **GPUs (all regions)** and **NVIDIA L4 GPUs** (regional) - must be ≥ 1. -2. **Kaggle Credentials**: If you are using gated models like Gemma 3, - you need a Kaggle account with the - [model license accepted](https://www.kaggle.com/models/keras/gemma3), - and `KAGGLE_USERNAME` / `KAGGLE_KEY` set in your local environment. +1. **A node pool** for your accelerator (default scale-to-zero). Pick any + slice/GPU that fits your model: -## Configuration + ```bash + kinetic pool add --accelerator tpu-v5litepod-1 --project your-project-id # TPU + kinetic pool add --accelerator gpu-l4 --project your-project-id # GPU + ``` -To run vLLM successfully on GPU via Kinetic, you need to handle -dependencies and environment variables properly. + 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`. -### 1. Dependencies - -Create a `requirements.txt` file in the directory of your script -containing: - -```text -keras -keras-hub -tensorflow-text -vllm -``` +2. **A `requirements.txt`** next to the scripts. The base set depends on the + device; a non-default export backend adds one entry: -Kinetic will detect this file and build a container with vLLM installed. -`tensorflow-text` is required — -KerasHub tokenizers preprocess with `tf.data` on every backend (CPU-side -only). Use a **Python 3.12** local venv; the remote container matches your -local interpreter, and `tensorflow-text` doesn't publish wheels for the -newest Python yet. + | 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]` | -### 2. Environment Variables + 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). -The example sets the following environment variables on the remote worker -to ensure correct execution: +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`. -- `KERAS_BACKEND="torch"`: vLLM is PyTorch-based, and the torch backend - is the only one that releases VRAM cleanly after the export, so vLLM's - KV cache gets the full GPU. -- `VLLM_USE_FLASHINFER_SAMPLER="0"`: vLLM's FlashInfer sampler - JIT-compiles CUDA kernels with `nvcc`, which pip-only containers don't - have; this selects the native torch sampler instead. -- `LD_LIBRARY_PATH=/usr/local/nvidia/lib64:...`: GKE mounts the host - NVIDIA driver at `/usr/local/nvidia`; this makes `libcuda` / - `libnvidia-ml` visible to vLLM's spawned engine processes (the example - also preloads them into the main process via `ctypes`). +## The example -These are set inside the remote function by the example itself. Kaggle -credentials are forwarded from your local environment via -`capture_env_vars` in the `@kinetic.run` decorator. +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). -## Example +```{literalinclude} ../../examples/export_worker.py +:language: python +:caption: examples/export_worker.py +``` ```{literalinclude} ../../examples/vllm_serving.py :language: python +:caption: examples/vllm_serving.py ``` -## Running the Example +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 -python3 vllm_serving.py +python vllm_serving.py ``` -The first run builds the container image (15–25 minutes; subsequent runs -reuse it as a cache hit) and provisions a GPU node from the scale-to-zero -pool (~10 minutes including the image pull). Monitor from a second -terminal with `kinetic jobs list` and +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`. -## Serving on TPU - -The export step is accelerator-agnostic, but vLLM serving itself always -needs a GPU or TPU. TPU serving uses a different vLLM build (`vllm-tpu`) -and different environment variables than the GPU example above, and -Kinetic builds one container per script directory — `vllm` and `vllm-tpu` -cannot share an image. So on TPU, export and serving run as **two scripts -in two directories, each with its own `requirements.txt`**. The checkpoint -moves between them through GCS; nothing is downloaded to your local -machine, and no re-initialization of Kinetic is needed. - -1. **Export script** (directory A, `requirements.txt`: `keras`, - `keras-hub`, `tensorflow-text`) — the export half of this example, on - any accelerator (`"cpu"` works). It uploads the checkpoint directory to - `KINETIC_OUTPUT_DIR` (a GCS path), which it returns; pass that path to - the serving script. -2. **Serving script** (directory B, `requirements.txt`: `vllm-tpu`) — - configured per - [Running vLLM on TPU with Kinetic](../guides/vllm_tpu.md): set - `VLLM_TARGET_DEVICE="tpu"`, `VLLM_USE_V1="0"`, and - `JAX_PLATFORMS="tpu,cpu"` locally and forward them with - `capture_env_vars`. The job downloads the checkpoint directory from - GCS to the pod's local disk and points `LLM(model=...)` at it instead - of a Hub model ID: - -```python -@kinetic.run( - accelerator="tpu-v5litepod-8", - capture_env_vars=["VLLM_*", "JAX_*"], -) -def serve_on_tpu(checkpoint_gs_path, prompts): - # Download the exported checkpoint directory from GCS to - # /tmp/hf_export (google.cloud.storage transfer_manager), then: - from vllm import LLM, SamplingParams - - llm = LLM(model="/tmp/hf_export", max_model_len=1024) - return llm.generate(prompts, SamplingParams(max_tokens=128)) -``` +## 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/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 index 78cbbd30..5b4de388 100644 --- a/examples/vllm_serving.py +++ b/examples/vllm_serving.py @@ -1,63 +1,30 @@ -"""Export a KerasHub causal LM to the Hugging Face Transformers format -and serve it with vLLM on a Kinetic GPU, in a single job. - -See the accompanying guide for prerequisites and configuration. -""" +"""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 -# Any preset with a Transformers exporter works (Gemma 3 / Gemma / -# Qwen / GPT-2). +DEVICE = "tpu" # "tpu" | "gpu" +BACKEND = ( + "jax" if DEVICE == "tpu" else "torch" +) # "jax" | "torch" | "tensorflow" MODEL_PRESET = "gemma3_4b" - -# Load and export in the model's native precision; "float32" for GPT-2 -# or pre-Ampere GPUs. DTYPE = "bfloat16" - EXPORT_DIR = "/tmp/hf_export" -# vLLM routes checkpoints via the `architectures` key in config.json. -# The Gemma 3 exporter writes it natively; the others don't yet, so we -# patch it after export. TODO: drop once all exporters in -# `keras_hub/src/utils/transformers/export/` write this key. -HF_ARCHITECTURES = { - "gpt2": "GPT2LMHeadModel", - "qwen2": "Qwen2ForCausalLM", - "gemma": "GemmaForCausalLM", - "gemma3_text": "Gemma3ForCausalLM", -} - - -def _ensure_architectures(export_dir): - """Add the `architectures` key to config.json if missing.""" - import json - - config_path = os.path.join(export_dir, "config.json") - with open(config_path) as f: - config = json.load(f) - if "architectures" not in config: - arch = HF_ARCHITECTURES.get(config.get("model_type")) - if arch is None: - raise ValueError( - f"Unknown model_type {config.get('model_type')!r}; " - "add it to HF_ARCHITECTURES." - ) - config["architectures"] = [arch] - with open(config_path, "w") as f: - json.dump(config, f, indent=2) +# 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-mounted NVIDIA driver to torch/vLLM and select - the torch-native sampler (pip-only images have no nvcc).""" + # Expose the GKE NVIDIA driver so the export child (which inherits + # LD_LIBRARY_PATH) and vLLM can use the GPU. import ctypes - import glob - - print("GPU devices:", glob.glob("/dev/nvidia*")) - # For vLLM's spawned engine processes, which read this at startup. nvidia_dirs = [ d for d in ("/usr/local/nvidia/lib64", "/usr/local/nvidia/lib") @@ -68,129 +35,65 @@ def _setup_gpu_runtime(): os.environ["LD_LIBRARY_PATH"] = ":".join( nvidia_dirs + ([prev] if prev else []) ) - - # For this process, whose linker ignores late LD_LIBRARY_PATH edits. for lib in ("libcuda.so.1", "libnvidia-ml.so.1"): - for root in ("/usr/local/nvidia/lib64", "/usr/local/nvidia/lib"): - path = os.path.join(root, lib) - if os.path.exists(path): - ctypes.CDLL(path, mode=ctypes.RTLD_GLOBAL) - print("preloaded:", path) + 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 - else: - print("NOT FOUND:", lib) - os.environ["VLLM_USE_FLASHINFER_SAMPLER"] = "0" -def _persist_export(export_dir): - """Upload the exported checkpoint directory to the job's durable - output dir.""" - import shutil - - output_dir = os.environ.get("KINETIC_OUTPUT_DIR") - if not output_dir: - return None - dest = f"{output_dir.rstrip('/')}/hf_export" - if dest.startswith("gs://"): - # Direct directory upload: no tar copy on ephemeral disk, and - # parallel workers make it fast. - from google.cloud import storage - from google.cloud.storage import transfer_manager - - bucket_name, _, blob_prefix = dest[5:].partition("/") - bucket = storage.Client().bucket(bucket_name) - files = [] - for root, _, filenames in os.walk(export_dir): - for filename in filenames: - rel_path = os.path.relpath(os.path.join(root, filename), export_dir) - files.append(rel_path) - print("Uploading exported checkpoint to GCS...") - transfer_manager.upload_many_from_filenames( - bucket, - files, - source_directory=export_dir, - blob_name_prefix=blob_prefix + "/", - worker_type=transfer_manager.THREAD, - raise_exception=True, - ) - else: - shutil.copytree(export_dir, dest, dirs_exist_ok=True) - return dest - - @kinetic.run( - accelerator="gpu-l4", + accelerator=ACCELERATOR, capture_env_vars=["KAGGLE_*", "GOOGLE_CLOUD_*"], ) def export_and_serve(prompts): - _setup_gpu_runtime() - - # torch is the only backend that frees VRAM for vLLM in the same - # process; keras reads this at first (transitive) import. - os.environ["KERAS_BACKEND"] = "torch" - - import gc - - import keras_hub - import torch - - # Forwarded via capture_env_vars; ~/.kaggle/kaggle.json doesn't - # travel to the pod. - print( - "Kaggle creds in pod:", - bool(os.environ.get("KAGGLE_USERNAME")) - and bool(os.environ.get("KAGGLE_KEY")), - ) - - print(f"Loading {MODEL_PRESET} from KerasHub ({DTYPE})...") - lm = keras_hub.models.CausalLM.from_preset(MODEL_PRESET, dtype=DTYPE) - - print("Exporting to Hugging Face Transformers format...") - lm.export_to_transformers(EXPORT_DIR) - _ensure_architectures(EXPORT_DIR) - - # Persist before starting vLLM: engine init is the riskiest step, - # and the checkpoint survives even if it crashes. - artifact = _persist_export(EXPORT_DIR) - if artifact: - print(f"Exported model uploaded to: {artifact}") - - # Release VRAM before handing the GPU to vLLM. - del lm - gc.collect() - torch.cuda.empty_cache() + 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 - print("Starting vLLM engine from the exported checkpoint...") llm = LLM( - model=EXPORT_DIR, - dtype=DTYPE, - gpu_memory_utilization=0.85, - max_model_len=1024, - enforce_eager=False, + model=EXPORT_DIR, load_format="safetensors", dtype=DTYPE, **serve_kwargs ) - sampling_params = SamplingParams(temperature=0.6, top_p=0.9, max_tokens=128) - - # All prompts are processed in one continuous-batching pass. - outputs = llm.generate(prompts, sampling_params) - results = [ + 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 ] - return results - 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", - "The most underrated skill in software engineering is", ] - for result in export_and_serve(prompts): + for r in export_and_serve(prompts): print("=" * 60) - print(f"Prompt: {result['prompt']}") - print(f"Completion: {result['completion']}") + print("Prompt:", r["prompt"]) + print("Completion:", r["completion"])