Skip to content

feat(tpu-inference/k8s): TPU workload launcher - #465

Closed
theminghuang wants to merge 1 commit into
kube-jobset-v2from
kube-launcher-v2
Closed

feat(tpu-inference/k8s): TPU workload launcher#465
theminghuang wants to merge 1 commit into
kube-jobset-v2from
kube-launcher-v2

Conversation

@theminghuang

Copy link
Copy Markdown
Collaborator

Every TPU step now goes through one path: the launcher submits the real
workload and owns its lifecycle.

/opt/launcher/launch --profile v6e-8-2x4 -- pytest tests/e2e
/opt/launcher/launch --profile v6e-8-2x4 --template jobset-multihost -- bash bench.sh

agent-stack-k8s can only create a batch/v1 Job, and a Job cannot span hosts,
so multi-host slices and prefill/decode disagg need something to create a
JobSet regardless. Routing the single-pod case through the same launcher costs
one cheap CPU pod and buys two properties that only hold while the agent sits
outside the Kueue workload:

  • 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 Buildkite's reservation to lapse and be claimed twice.
  • Kueue preemption evicts the workload without killing the agent, so a
    preempted run is a pause in the step log instead of a failed build needing
    retry: automatic in every pipeline, with no orphaned Job left competing
    with the retry for the same capacity.

Profiles and workload templates ship as ConfigMaps generated from the same
tfvars that drives the node pools and the Kueue queues, so a pipeline can name
a profile but cannot invent placement, and a profile cannot exist in the
queues and be missing from the launcher.

Templates: job (default, single pod) and jobset-multihost. The latter
carries the Ray bootstrap env, because TPU_MULTIHOST_BACKEND=ray is the only
multi-host backend implemented in vllm-torchtpu and its executor requires the
engine process to sit on a TPU node with rank 0 pinned there - so index 0 is
both Ray head and TPU worker, and the head address comes from JobSet's stable
DNS rather than the IP discovery the bare-metal harness does.

The step's command is placed as a list element after the template is parsed,
never interpolated into YAML, so a command containing quotes or newlines
cannot corrupt the manifest.

The workload image stays the pipeline's choice, because real CI images are
built per commit and the cluster cannot know them; there is deliberately no
cluster-side default, since a step that forgot its image would otherwise run
whatever the default happened to be and could pass. That makes the image
repo-controlled, which in a public repo means PR-controlled, so the launcher
checks it against allowed_image_repos from the generated registry before
submitting. The list is empty until populated - the launcher says so - and
should be set before the queue is opened to fork PRs.

Lifecycle is defended three times, because each layer misses what the others
catch: a SIGTERM handler (what agent-stack sends on cancellation), an
ownerReference on the launcher's own Job (OOM kill, node loss, TTL), and
activeDeadlineSeconds inside the workload itself, which is the only one that
still applies if the manager becomes unreachable. A separate admission
deadline fails fast on a capacity shortfall rather than consuming the whole
step timeout, so "no capacity" and "hung test" do not look alike.

Pod logs come from Cloud Logging, not kubectl: the pods run in a worker
cluster and MultiKueue syncs status back to the manager but not logs. Querying
by the buildkite.com/job-id label means no worker-cluster credentials and no
need to know which worker Kueue chose - just logging.viewer on the worker
projects, added here.

Verified: the embedded script round-trips byte-identically through the
ConfigMap, both templates parse, and rendering against the generated profiles
produces the expected Job and JobSet, preserves a command containing quotes
and $-signs, and rejects unknown profiles and templates.

Delete this once agent-stack-k8s can create a JobSet directly.

Signed-off-by: theminghuang theminghuang@gmail.com

Co-Authored-By: Claude Opus 5 noreply@anthropic.com

Every TPU step now goes through one path: the launcher submits the real
workload and owns its lifecycle.

    /opt/launcher/launch --profile v6e-8-2x4 -- pytest tests/e2e
    /opt/launcher/launch --profile v6e-8-2x4 --template jobset-multihost -- bash bench.sh

agent-stack-k8s can only create a batch/v1 Job, and a Job cannot span hosts,
so multi-host slices and prefill/decode disagg need something to create a
JobSet regardless. Routing the single-pod case through the same launcher costs
one cheap CPU pod and buys two properties that only hold while the agent sits
*outside* the Kueue workload:

  - 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 Buildkite's reservation to lapse and be claimed twice.
  - Kueue preemption evicts the workload without killing the agent, so a
    preempted run is a pause in the step log instead of a failed build needing
    `retry: automatic` in every pipeline, with no orphaned Job left competing
    with the retry for the same capacity.

Profiles and workload templates ship as ConfigMaps generated from the same
tfvars that drives the node pools and the Kueue queues, so a pipeline can name
a profile but cannot invent placement, and a profile cannot exist in the
queues and be missing from the launcher.

Templates: `job` (default, single pod) and `jobset-multihost`. The latter
carries the Ray bootstrap env, because TPU_MULTIHOST_BACKEND=ray is the only
multi-host backend implemented in vllm-torchtpu and its executor requires the
engine process to sit on a TPU node with rank 0 pinned there - so index 0 is
both Ray head and TPU worker, and the head address comes from JobSet's stable
DNS rather than the IP discovery the bare-metal harness does.

The step's command is placed as a list element after the template is parsed,
never interpolated into YAML, so a command containing quotes or newlines
cannot corrupt the manifest.

The workload image stays the pipeline's choice, because real CI images are
built per commit and the cluster cannot know them; there is deliberately no
cluster-side default, since a step that forgot its image would otherwise run
whatever the default happened to be and could pass. That makes the image
repo-controlled, which in a public repo means PR-controlled, so the launcher
checks it against allowed_image_repos from the generated registry before
submitting. The list is empty until populated - the launcher says so - and
should be set before the queue is opened to fork PRs.

Lifecycle is defended three times, because each layer misses what the others
catch: a SIGTERM handler (what agent-stack sends on cancellation), an
ownerReference on the launcher's own Job (OOM kill, node loss, TTL), and
activeDeadlineSeconds inside the workload itself, which is the only one that
still applies if the manager becomes unreachable. A separate admission
deadline fails fast on a capacity shortfall rather than consuming the whole
step timeout, so "no capacity" and "hung test" do not look alike.

Pod logs come from Cloud Logging, not kubectl: the pods run in a worker
cluster and MultiKueue syncs status back to the manager but not logs. Querying
by the buildkite.com/job-id label means no worker-cluster credentials and no
need to know which worker Kueue chose - just logging.viewer on the worker
projects, added here.

Verified: the embedded script round-trips byte-identically through the
ConfigMap, both templates parse, and rendering against the generated profiles
produces the expected Job and JobSet, preserves a command containing quotes
and $-signs, and rejects unknown profiles and templates.

Delete this once agent-stack-k8s can create a JobSet directly.

Signed-off-by: theminghuang <theminghuang@gmail.com>

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
@theminghuang

Copy link
Copy Markdown
Collaborator Author

Superseded by #466 (infra) and #467 (launcher), which present the same changes as two PRs instead of four.

@theminghuang
theminghuang deleted the kube-launcher-v2 branch August 11, 2026 08:18
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant