Add guide: Serving KerasHub models with vLLM using Kinetic

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.

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::

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)

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"])