Add ai-ml/gke-ray/tpu/get-started: Ray on TPU end-to-end example#2122
Merged
olivi-eh merged 9 commits intoJul 17, 2026
Merged
Conversation
End-to-end Ray-on-TPU get-started example on a single Cloud TPU v6e slice: - cluster: GKE + Ray Operator add-on + v6e slice + GCS + monitoring (Terraform) - serve: Qwen3-4B behind an OpenAI-compatible API (Ray Serve + vLLM) - data: DPO dataset prep + batch prediction (Ray Data) - train: DPO fine-tune with LoRA (Ray Train + Tunix) Adds a CI workflow that validates the Terraform and builds the serve and train images. Towards GoogleCloudPlatform#2119.
Replace the Step/Content table in the get-started README with a numbered list and a short lead-in, so the end-to-end order reads more clearly.
Expand each step in the get-started README to say what it does and which Ray library it uses, and reference each sample directory inline.
Replace the loose pip specs in the serve and train Dockerfiles with the exact versions frozen from a build on the pinned vllm-tpu base, for reproducibility: - serve: protobuf==5.29.6 (grpcio-status stays 1.70.0) - train: protobuf==6.33.6, google-tunix==0.1.7, grain==0.2.18, qwix==0.1.8, wandb==0.28.0
The provider "helm" block uses the v3 `kubernetes = {}` attribute syntax, so
helm 2.x is incompatible, not just low.
- google: >= 6.0.0 -> >= 7.0.0
- kubernetes: >= 2.30.0 -> >= 3.0.0
- helm: >= 2.13.0 -> >= 3.0.0
terraform validate passes with the bumped floors.
topology and accelerator_type are already CLI args; add a matching --num-workers (default 1) so the trainer scales to multi-host slices. Default preserves the validated single-host 2x4 behavior.
Add the Ray 2.56 high-throughput env vars to the serve RayService (head and worker) and install haproxy in the serve image, which they require. Document them in the serve README. Verified on a v6e slice: proxies HEALTHY, HAProxy actors ALIVE, endpoint serves.
Add a prerequisites lead-in to the cluster and train READMEs.
Add a Prerequisites section to the data README (which Ray head the jobs use; batch prediction needs the train RayCluster for /data + TPU), and note the trainer's --topology/--num-workers flags for multi-host slices.
inardini
requested review from
a team,
alizaidis,
moficodes and
yoshi-approver
as code owners
July 14, 2026 16:57
|
Here is the summary of changes. You are about to add 7 region tags.
This comment is generated by snippet-bot.
|
andrewsykim
approved these changes
Jul 16, 2026
ryanaoleary
approved these changes
Jul 16, 2026
ryanaoleary
left a comment
ryanaoleary
left a comment
Contributor
There was a problem hiding this comment.
Primarily reviewed the python scripts, YAML, and READMEs but LGTM!!
spencer-p
approved these changes
Jul 17, 2026
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
6 participants
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.
Description
Adds a new end-to-end Ray-on-TPU get-started example under
ai-ml/gke-ray/tpu/get-started/, alongside the existingtpu/samples. It walksthrough the full LLM lifecycle on a single Cloud TPU v6e slice with Ray on GKE,
using Qwen3-4B-Instruct-2507 (small, text-only, Apache-2.0) and the
ultrafeedback_binarizedpreference dataset:single-host v6e
2x4slice (Spot), Workload Identity, a Cloud Storage bucket,and a Prometheus + Grafana monitoring stack.
including the Ray 2.56 high-throughput ingress (HAProxy).
across the slice.
Proposed and accepted in #2119.
Notes for reviewers:
vllm-tpunightly base (captured from apip freezeof the built images). Theserve and train images build via Cloud Build.
2x4slice, including the high-throughput servepath (proxies healthy, HAProxy actors alive, endpoint serves).
vLLM, and
ray.serve.llmon TPU. The serve and train READMEs document this andrecommend mirroring it to your own registry for durability.
Tasks