tpu-inference/k8s: MultiKueue GKE infrastructure for TPU CI

.gitignore

@@ -15,6 +15,7 @@ crash.*.log
 # to change depending on the environment.
 *.tfvars
 *.tfvars.json
+!terraform/gcp_old/tpu-inference/k8s/prod.auto.tfvars
 
 # Ignore override files as they are usually used to override resources locally and so
 # are not checked in

terraform/gcp_old/tpu-inference/k8s/README.md

@@ -0,0 +1,166 @@
+# MultiKueue TPU CI/CD Infrastructure on GKE
+
+This directory contains the Terraform configuration, Kueue manifest templates, and manifest generator script for managing the MultiKueue-based TPU CI/CD testing infrastructure across GKE clusters.
+
+---
+
+## 1. Overview & Architecture
+
+The infrastructure uses a **MultiKueue** architecture to distribute TPU CI benchmark workloads submitted via Buildkite across dedicated GKE clusters in Google Cloud Platform:
+
+```
+                  +----------------------------------------------+
+                  |              Manager Cluster                 |
+                  |     Project: cloud-ullm-inference-ci-cd      |
+                  |           Name: tpu-ci-manager               |
+                  |          Region: us-central1                 |
+                  +----------------------+-----------------------+
+                                         |
+                                Buildkite Agent Stack
+                                    (Controller)
+                                         |
+                                 Kueue MultiKueue
+                               (Admission & Dispatch)
+                                         |
+                                         v
+                     +----------------------------------------+
+                     |             Worker Cluster             |
+                     |      Project: cloud-tpu-inference-test |
+                     | Name: tpu-ci-southamerica-west1-a      |
+                     |     Location: southamerica-west1-a     |
+                     +-------------------+--------------------+
+                                         |
+                                 GKE TPU Node Pools
+                             (v6e-1: 1-chip, v6e-8: 8-chip)
+```
+
+---
+
+## 2. Key Accomplishments & Design Principles
+
+### A. Centralized Buildkite Controller
+To prevent race conditions where worker clusters compete to claim the same Buildkite job, the `agent-stack-k8s` controller runs **exclusively on the manager cluster**, and there is exactly **one** of it, on a single queue.
+
+The Buildkite queue deliberately carries no TPU information. A step names its shape with `--profile`, and the controller stays shape-agnostic, so adding a TPU profile is a regenerated ConfigMap rather than another Helm release. The controller's own pod is CPU-only and carries no Kueue queue label, so Kueue never queues or evicts it - only the workload it submits.
+
+### B. MultiKueue Dispatch over Connect Gateway
+Kueue on the manager cluster inspects submitted jobs, matches cluster queues (`v6e-1-1x1`, `v6e-8-2x4`), and dispatches the workload object to the corresponding worker cluster (`southamerica-west1-a`) via GKE Connect Gateway using Workload Identity (`roles/gkehub.gatewayEditor`).
+
+### C. Native Cross-Project Image Pulling (No K8s Secrets Required)
+Container images for testing (such as `us-central1-docker.pkg.dev/cloud-ullm-inference-ci-cd/tpu-inference-ci/vllm-tpu`) are hosted in the manager project. 
+
+Terraform codifies cross-project IAM reader permissions (`roles/artifactregistry.reader`) for worker node service accounts (`tpu-ci-wkr-node@cloud-tpu-inference-test.iam.gserviceaccount.com`) on the manager project. GKE node containerd runtimes authenticate natively via GCP metadata tokens without requiring manual Kubernetes `imagePullSecrets` or service account keys.
+
+### D. TPU is the only resource under quota
+Buildkite workloads generate helper containers (e.g. `copy-agent`), initContainers (`imagecheck`) and the gcsfuse sidecar, all of which request CPU and memory. Kueue's default `quotaCheckStrategy: BlockUndeclared` refuses to admit a workload that requests a resource its ClusterQueue does not cover, so an earlier version of the queues carried `cpu` and `memory` in `coveredResources` with quotas of `10000` / `10000Gi` as a stand-in for "unbounded" - a `nominalQuota` of `0` is a real zero, not unlimited.
+
+Both controller configs now set `resources.quotaCheckStrategy: IgnoreUndeclared` (Kueue 0.19, feature gate on by default), under which only the resources a ClusterQueue declares are checked. `queue_group.yaml.tpl` covers `google.com/tpu` alone; CPU and memory are enforced by the kube scheduler against node capacity, which is where they belong. Manager and workers must carry the same setting, since the worker's Kueue admits the mirrored workload too.
+
+### E. Standardized Node Pool Configuration
+TPU node pools in `clusters.tf` are configured with:
+- **Taints**: `google.com/tpu=present:NoSchedule` (allows GKE Cluster Autoscaler to simulate scale-up for pending TPU pods).
+- **Reservation Affinity**: `reservation_affinity` targeting designated Cloud TPU reservations.
+- **Placement by workload, not by flavor**: the `ResourceFlavor` is per-accelerator (`v6e`) and carries no `nodeLabels`, so every profile in the cohort shares one flavor and can borrow from it. Kueue therefore injects no node selector, and the submitted workload names `accelerator_label` and `topology` itself. That is what makes a 1-chip job and an 8-chip job draw on the same pool of chips.
+
+### F. JobSet for multi-pod workloads
+A `batch/v1` Job cannot span hosts, so multi-host slices and prefill/decode disaggregation need JobSet. The operator is installed on the manager and every worker from a single `jobset_version`, and `jobset.x-k8s.io/jobset` is enabled in `integrations.frameworks` on both sides - MultiKueue mirrors the object across clusters, so the CRD, the operator version and the enabled framework list all have to match.
+
+A pool with `hosts > 1` is a multi-host TPU slice pool: GKE creates it with a `COMPACT` placement policy carrying the topology, one node per host, and the autoscaler scales it atomically - every host or none. `clusters.tf` checks that such a pool's node counts are whole slices and that one pool is one slice (`max_nodes == hosts`); more slices of a shape are more pools. A four-host `ct6e-standard-4t` 4x4 slice (16 chips, a JobSet with `parallelism 4`) has been run this way; none is declared in the tfvars today.
+
+### G. Physical shapes constrain borrowing
+Node pools are single-shape, so with an 18-chip reservation one 8-chip node and ten 1-chip nodes leave nothing for a 16-chip slice. Quota borrowing across shapes is therefore not free: reclaiming chips means draining and deleting nodes of one shape before nodes of the other can be created. Measured on this cluster, that costs roughly 300s on top of the ~110s node scale-up. Worth knowing before tuning quotas - the cost is physical, not a Kueue setting.
+
+---
+
+## 3. Workflow & Usage
+
+### Modifying Configuration
+All cluster topology and pool limits are declared in `prod.auto.tfvars`.
+
+Example pool definition:
+```hcl
+tpu_pools = {
+  v6e-1-1x1 = {
+    machine_type      = "ct6e-standard-1t"
+    accelerator       = "v6e"          # names the cohort and the ResourceFlavor
+    accelerator_label = "tpu-v6e-slice"
+    topology          = "1x1"
+    chips_per_node    = 1
+    min_nodes         = 2              # kept warm
+    nominal_nodes     = 10             # Kueue nominalQuota = chips x nominal_nodes
+    max_nodes         = 10             # autoscaling ceiling; anything above nominal is borrowed
+    reservation_name  = "cloudtpu-20250327121505-861300654"
+  }
+  # a slice across hosts would add `hosts = 4` on a ct6e-standard-4t 4x4 pool;
+  # one multi-host pool is one slice and max_nodes must equal hosts.
+}
+```
+
+The pool key is the profile name, and it is used verbatim as the Kueue
+ClusterQueue and LocalQueue name and as the `--profile` a pipeline passes.
+`nominal_nodes` is what the profile owns; anything between it and `max_nodes`
+is borrowed from the cohort and is the first thing reclaimed when another
+profile needs its own quota back.
+
+### Manifest Generation
+
+```bash
+python3 -m pip install -r scripts/requirements.txt   # hcl2, once
+python3 scripts/generate_manifests.py
+```
+
+Manifests are rendered into per-cluster directories, numbered in the order
+`kubectl` applies them:
+
+```
+generated/manager/            01-base 02-multikueue-fleet 03-cohorts
+                              04-resource-flavors 05-queues [06-launcher]
+generated/worker-<location>/  01-base 02-resource-flavors 03-queues
+                              [04-launcher-rbac]
+```
+
+Only `01-base` ordering is load-bearing - it carries the Namespace everything
+else is created into. The rest is soft: a ClusterQueue naming a ResourceFlavor
+that does not exist yet goes inactive and recovers when it appears.
+
+`generated/` is committed, so regenerating should produce no diff unless you
+changed a template or the tfvars. A non-empty diff after an unrelated change
+means something drifted.
+
+### Applying Infrastructure & Manifests
+
+```bash
+terraform fmt && terraform apply
+./scripts/deploy_manifests.sh
+```
+
+`deploy_manifests.sh` applies each cluster's directory in one call. To land
+them one at a time when debugging a fresh install:
+
+```bash
+for f in generated/manager/*.yaml; do echo "== $f"; kubectl apply -f "$f" || break; done
+```
+
+### Verifying a deploy
+
+Against the **manager**:
+
+```bash
+kubectl -n buildkite get localqueue
+kubectl get crd jobsets.jobset.x-k8s.io
+kubectl -n buildkite get pods -l app.kubernetes.io/name=agent-stack-k8s
+```
+
+Expect LocalQueues matching the profile names, the JobSet CRD, and exactly one
+agent-stack pod - more than one means an older per-profile controller survived
+and will compete for jobs.
+
+---
+
+## 4. Best Practices & Troubleshooting
+
+### Build Cancellation
+Always cancel builds via the **Buildkite UI** or Buildkite CLI (`buildkite-agent build cancel <build-id>`). Avoid running manual `kubectl delete job` directly out-of-band, as deleting `Job` objects bypasses the controller event watcher and leaves Buildkite builds pending until step timeouts expire.
+
+### Pipeline Timeouts
+In `.buildkite/pipeline_kube.yaml`, ensure pipeline steps include `timeout_in_minutes: N` and `cancel_on_build_failing: true` so stalled steps fail fast automatically.

terraform/gcp_old/tpu-inference/k8s/backend.tf

@@ -0,0 +1,6 @@
+terraform {
+  backend "gcs" {
+    bucket = "cloud-ullm-inference-ci-cd-tf-state"
+    prefix = "buildkite-tpu-ci/foundation"
+  }
+}

terraform/gcp_old/tpu-inference/k8s/buildkite.tf

@@ -0,0 +1,70 @@
+# Buildkite Agent Stack (manager cluster only, so worker clusters never race
+# to claim Buildkite jobs).
+#
+# One controller, one queue. Every TPU step goes through the launcher, which
+# submits the real workload as a Kueue-managed object, so the Buildkite queue
+# no longer has to encode a TPU shape - the profile is an argument to `launch`.
+# Adding a TPU profile is then a regenerated ConfigMap rather than another Helm
+# release.
+#
+# The agent pod is CPU-only and carries no Kueue queue label, so Kueue never
+# queues or evicts it. Two properties follow:
+#
+#   - The agent acquires its Buildkite job in seconds rather than after TPU
+#     admission and node pool scale-up, so the job is never held reserved long
+#     enough for the reservation to lapse and be picked up twice.
+#   - Kueue preemption evicts the workload without killing the agent, so a
+#     preempted run is a pause in the step log rather than a failed build
+#     needing `retry: automatic`.
+#
+# Placement - node selector, TPU resources, Kueue queue label - belongs to the
+# submitted workload, not here.
+
+resource "kubernetes_namespace_v1" "buildkite_manager" {
+  provider = kubernetes.manager
+
+  metadata {
+    name = "buildkite"
+    labels = {
+      "pod-security.kubernetes.io/enforce" = "baseline"
+      "pod-security.kubernetes.io/audit"   = "restricted"
+      "pod-security.kubernetes.io/warn"    = "restricted"
+    }
+  }
+}
+
+resource "helm_release" "buildkite_agent_stack" {
+  provider         = helm.manager
+  name             = "agent-stack-k8s"
+  repository       = "oci://ghcr.io/buildkite/helm"
+  chart            = "agent-stack-k8s"
+  version          = var.buildkite_agent_stack_chart_version
+  namespace        = kubernetes_namespace_v1.buildkite_manager.metadata[0].name
+  create_namespace = false
+  force_update     = true
+  cleanup_on_fail  = true
+  replace          = true
+
+  values = [
+    yamlencode({
+      agentStackSecret = "agent-stack-k8s-secret"
+
+      config = {
+        id    = "tpu-ci"
+        queue = var.buildkite_queue
+        debug = var.buildkite_agent_stack_debug
+
+        # Applies to the launcher pod only; the workload carries its own
+        # activeDeadlineSeconds. Sized to outlast queue wait plus the run,
+        # because the launcher waits for both.
+        job-active-deadline-seconds = var.tpu_job_max_runtime_seconds
+        pod-pending-timeout         = "180m"
+      }
+    })
+  ]
+
+  depends_on = [
+    kubernetes_namespace_v1.buildkite_manager,
+    null_resource.manager_external_secrets_helm
+  ]
+}