# Aiml On Kubernetes

> AI/ML on Kubernetes & GKE

- Skill: `sanjeevrg89/aiml-on-kubernetes` (Agent Skill, multi-file: 5 files)
- Install (CLI): `npx skillmds@latest add sanjeevrg89/aiml-on-kubernetes`
- Raw SKILL.md: https://api.skillmd.com/api/skills/sanjeevrg89/aiml-on-kubernetes/raw
- Safety review: pending
- Works with: Claude Code, Claude.ai, OpenAI Codex
- Category: DevOps & Infra
- Author: sanjeevrg89 (https://skillmd.com/u/sanjeevrg89)
- Updated: 2026-09-17
- Page: https://skillmd.com/skills/sanjeevrg89/aiml-on-kubernetes

---


# AI/ML on Kubernetes & GKE

Apply the judgment of a staff ML-infra engineer who runs frontier training and serving on Kubernetes and
GKE: you size accelerators, schedule gang jobs that either run whole or not at all, wire up RDMA/TCPX so
NCCL hits line rate, keep MFU and goodput high across thousands of chips, and serve trillion-parameter
models across hosts. This is the **map of the territory** — it gives the end-to-end lifecycle and the
K8s/GKE orchestration, and routes to the deeper sibling skills for framework internals.

## How to use this skill

1. **Read `aiml-on-kubernetes-guide.md`** in this directory — the full reference: accelerators on
   GKE, distributed training, inference/serving, fine-tuning, RL/RLHF, agentic, cross-cutting infra,
   and a decision guide. Apply it to the task.
2. For annotated end-to-end manifests (multi-host JobSet+Kueue training; multi-host vLLM/LWS serving),
   read **`examples.md`** and imitate them.
3. **Route to the right sibling** rather than going deep here: this skill is the integrator. Match the
   existing cluster's conventions (node-pool labels, taints, queue names); apply the correctness/safety
   and topology rules regardless.

## Essentials (full detail in `aiml-on-kubernetes-guide.md`)

- **Accelerators are scheduled like any resource, but placement is everything.** Request
  `nvidia.com/gpu` or `google.com/tpu`; the real work is node pools, drivers, taints/tolerations, and
  **topology-aware placement** so a multi-host slice lands on one contiguous network domain. On GKE,
  let the GPU/TPU device plugins and the driver installer DaemonSet do the low-level work; verify
  current SKU/region availability — hardware moves fast.
- **Multi-host = gang scheduling, full stop.** A job spanning N hosts must get all N pods or none, or
  you burn accelerators idling. Use [[kueue-advanced]] for queueing/quota/all-or-nothing admission and
  [[jobset-leaderworkerset]] (JobSet for training, LeaderWorkerSet for inference) to model the
  multi-pod group with stable identity and headless networking.
- **Networking makes or breaks training.** GPUDirect-TCPX/TCPXO or RDMA over Converged Ethernet, multi-NIC
  pods, and correct NCCL env are the difference between 30% and 90% of line rate. TPU multi-host slices
  use the ICI mesh — you don't manage it, but you must keep the slice contiguous (topology).
- **Checkpoint for fault tolerance, not just resumption.** At scale a host *will* fail mid-run; design
  for it. Asynchronous/multi-tier checkpointing to local disk → Hyperdisk/Parallelstore → GCS, plus
  elastic/restartable jobs. Goodput = useful compute / wall-clock; checkpoint overhead and restart cost
  are first-order.
- **Inference splits single-host vs multi-host.** If the model + KV cache fit one node, it's a Deployment
  with HPA-style autoscaling ([[autoscaling-kubernetes]]). If it doesn't, it's multi-host with
  LeaderWorkerSet, one logical replica = leader + workers sharing weights over the fabric.
- **Serving infra is shaped by KV cache and batching.** Continuous batching, prefix/KV-cache reuse, and
  prefill/decode disaggregation change your memory, autoscaling signal, and routing. Defer engine
  internals to [[serving-frameworks]]; own the GKE Inference Gateway / model-aware routing and load
  from GCS with fast cold-start.
- **Fine-tuning's infra ask depends on method.** Full fine-tune ≈ training (optimizer+grad sharding,
  multi-host). LoRA/QLoRA fits far smaller footprints — often single-node, sometimes a fraction of a
  GPU via MIG/time-slicing. Size the request to the method.
- **RL/RLHF is heterogeneous and multi-cluster by nature.** Rollout (inference) + reward model +
  learner (training) have different accelerator profiles and lifecycles; orchestrating their data flow
  and weight sync on K8s is the hard part. Name the frameworks (TRL, veRL, NeMo-RL, OpenRLHF); treat
  it as a systems problem.
- **Observability for accelerators is non-negotiable.** DCGM exporter for GPUs, TPU runtime metrics,
  GKE Managed Prometheus with namespace-level accelerator metrics. Watch MFU/goodput, SM/HBM occupancy,
  NCCL all-reduce time, XID/SXid errors, and host/NVLink failures.
- **Multi-tenancy is quota + isolation.** Kueue ClusterQueues/cohorts for fair-share of scarce
  accelerators; per-tenant namespaces, ResourceQuotas, and node taints. Reserved capacity (CUD/flex-start/
  DWS on GKE) vs on-demand changes how you queue.
- **AI FinOps: optimize price-per-result, not price-per-hour.** Accelerators are the dominant cost, so
  utilization *is* cost — the KPIs are MFU/goodput and $/training-run (training) and $/token &
  tokens/$/accelerator-hour at an SLO (serving). Levers: right-size SKUs, Spot/preemptible + checkpointing
  ([[ml-checkpointing-orbax]]), bin-pack (MIG/time-slicing), scale-to-zero ([[autoscaling-kubernetes]]),
  quota/fair-share ([[kueue-advanced]]), and the right reservation/on-demand/Spot/committed-use blend.
- **Plan capacity before you commit budget.** Estimate training accelerator-hours from
  ≈6·params·tokens / (peak-FLOPs·MFU); size inference from peak-QPS·tokens / measured tokens-sec-per-
  accelerator-at-SLO, plus headroom. Attribute cost by team via namespace-scoped accelerator metrics
  (`k8s_container`/`namespace_name`) → showback/chargeback and unit economics. See guide §10.

## Related skills

- `[[gke-master]]` — GKE Standard/Autopilot, TPU/GPU node pools, NAP, networking, security, DWS.
- `[[kueue-advanced]]` — batch queueing, quota/cohorts, gang admission, MultiKueue, Topology-Aware Scheduling.
- `[[jobset-leaderworkerset]]` — JobSet (multi-host training) and LeaderWorkerSet (multi-host inference).
- `[[ml-frameworks]]` — PyTorch, JAX, XLA, GPU/TPU programming internals.
- `[[training-frameworks]]` — DDP/FSDP, DeepSpeed/ZeRO, Megatron, NeMo, Ray Train, Kubeflow Trainer, MaxText.
- `[[serving-frameworks]]` — vLLM, SGLang, Dynamo, Triton, TensorRT-LLM, Ray Serve, KServe.
- `[[autoscaling-kubernetes]]` — HPA/VPA/Cluster Autoscaler/Karpenter/KEDA/NAP for model servers.
- `[[slurm-hpc-on-kubernetes]]` — Slurm/Slinky/Volcano/MPI/RDMA and the Slurm-vs-K8s decision.
- `[[kubernetes-expert]]` — general end-to-end Kubernetes practitioner mastery.