tpu-inference/k8s: the TPU workload launcher

[theminghuang]

The TPU workload launcher, on top of #438. Every TPU step in tpu-inference#3377 runs it:

/opt/launcher/launch --profile v6e-8-2x4 -- pytest tests/

What it does

In a CPU-only agent pod, the launcher renders the workload manifest that lives in the repo under test - a Job, or a JobSet for a slice across hosts - fills in chip count, topology, node labels and queue from the profile registry generated here, forwards the step's BUILDKITE_* environment plus a named list, submits it to the profile's Kueue queue, waits for admission for as long as the total budget allows less one full run, polls the pod logs back over Connect Gateway (short requests, not a stream; keyed by owning Job so a preempt-and-resume reads as one log), mirrors the container's exit code, and files the log as a Buildkite artifact. Cancellation is a SIGTERM that uploads and deletes the workload; an ownerReference from the launcher pod covers the cases that never deliver one.

Why

agent-stack-k8s can only create a plain batch/v1 Job, and the agent must stay outside the Kueue workload: a preempted or disrupted run then pauses the step log and resumes rather than failing the build, and an admission wait holds a CPU pod rather than a TPU. Routing single-host work through the same path keeps one code path.

The split of responsibility is the point: the profile registry is generated from the same tfvars that build the node pools, so placement cannot drift from the queues; the manifest lives in the repo under test, so the shape of a job is a PR rather than an infrastructure change.

The manifest is untrusted input

Validated before submission: the Kueue queue label must be the profile's (quota cannot be side-stepped), serviceAccountName must be one the cluster publishes for workloads (tpu-workload; a workload may not run as the launcher and submit further work), the image must come from allowed_image_repos (the CI registry, set in tfvars), and there must be a container named workload. PodSecurity baseline on the namespace covers privileged, hostPath and host networking.

Files

• kueue/launcher/launch.py - the launcher.
• kueue/templates/launcher.yaml.tpl -> generated/manager/06-launcher.yaml: ServiceAccount, Role, the launcher and the profile registry as ConfigMaps, and the PodTemplate agent-stack-k8s runs it in. launcher_rbac_worker.yaml.tpl -> 07-launcher-rbac.yaml: a per-worker ClusterRole for pod logs.
• scripts/generate_manifests.py - the launcher outputs; the profile registry (queue, chips, hosts, topology, accelerator label, max runtime, gcsfuse file-cache size per machine type); a check that a multi-host lane's quota is whole slices.
• Validated with a four-host JobSet slice on v6e (kube-dev build 272: TPU_WORKER_HOSTNAMES injected, one 16-device mesh, vLLM TP=16 served) as well as the single-host suite; the multi-host profile itself is not declared in this stack yet.
• iam.tf - Connect Gateway reader for the launcher identity on manager and workers; secret-scoped access to the Test Engine token.
• buildkite.tf, variables.tf - the agent deadline derived from tpu_total_max_seconds; the analytics-token secret ids.
• prod.auto.tfvars - allowed_image_repos.

[theminghuang]

This change is part of the following stack:

• tpu-inference/k8s: MultiKueue GKE infrastructure for TPU CI #438
  • tpu-inference/k8s: the TPU workload launcher #467 ◀

_{Change managed by git-spice.}