feat(tpu-inference/k8s): TPU workload launcher - #465
Closed
theminghuang wants to merge 1 commit into
Closed
Conversation
Collaborator
Author
|
This change is part of the following stack: Change managed by git-spice. |
This was referenced Aug 11, 2026
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
force-pushed
the
kube-jobset-v2
branch
from
August 11, 2026 08:06
52ace1e to
6fdbd38
Compare
theminghuang
force-pushed
the
kube-launcher-v2
branch
from
August 11, 2026 08:06
5e723c5 to
9de1c39
Compare
Collaborator
Author
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Every TPU step now goes through one path: the launcher submits the real
workload and owns its lifecycle.
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:
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.
preempted run is a pause in the step log instead of a failed build needing
retry: automaticin every pipeline, with no orphaned Job left competingwith 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) andjobset-multihost. The lattercarries 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