Skip to content

Commit 9860f5e

Browse files
ryanaolearyryanaoleary
authored
Add v6e KubeRay TPU samples (#1793)
* Add v6e KubeRay TPU samples Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> Add region tags and copyright Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> Update working_dir path Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> Fix yaml Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> Some fixes Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> Test change to tpu_list_devices Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> Try with placement group Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> Update samples Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> Add copyright Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> * Change paths to main Signed-off-by: Ryan O'Leary <ryanaoleary@google.com> --------- Signed-off-by: Ryan O'Leary <ryanaoleary@google.com>
1 parent 399f3b7 commit 9860f5e

3 files changed

Lines changed: 194 additions & 0 deletions

File tree

‎ai-ml/gke-ray/tpu/ray-job.tpu-v6e-multihost.yaml‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
# [START gke_ai_ml_gke_ray_tpu_ray_job_tpu_v6e_multihost]
16+
apiVersion: ray.io/v1
17+
kind: RayJob
18+
metadata:
19+
name: v6e-16-job
20+
spec:
21+
entrypoint: python ai-ml/gke-ray/tpu/tpu_list_devices.py
22+
runtimeEnvYAML: |
23+
working_dir: "https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/archive/refs/heads/main.zip"
24+
pip:
25+
- jax[tpu]==0.4.33
26+
- -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
27+
rayClusterSpec:
28+
rayVersion: '2.43.0'
29+
headGroupSpec:
30+
rayStartParams: {}
31+
template:
32+
spec:
33+
containers:
34+
- name: ray-head
35+
image: rayproject/ray:2.43.0-py310
36+
ports:
37+
- containerPort: 6379
38+
name: gcs-server
39+
- containerPort: 8265
40+
name: dashboard
41+
- containerPort: 10001
42+
name: client
43+
resources:
44+
limits:
45+
cpu: "8"
46+
memory: 40G
47+
requests:
48+
cpu: "8"
49+
memory: 40G
50+
workerGroupSpecs:
51+
- replicas: 1
52+
minReplicas: 1
53+
maxReplicas: 1
54+
numOfHosts: 4
55+
groupName: tpu-group
56+
rayStartParams: {}
57+
template:
58+
spec:
59+
containers:
60+
- name: ray-worker
61+
image: rayproject/ray:2.43.0-py310
62+
resources:
63+
limits:
64+
cpu: "24"
65+
google.com/tpu: "4"
66+
memory: 200G
67+
requests:
68+
cpu: "24"
69+
google.com/tpu: "4"
70+
memory: 200G
71+
env:
72+
- name: NODE_IP
73+
valueFrom:
74+
fieldRef:
75+
fieldPath: status.hostIP
76+
- name: VBAR_CONTROL_SERVICE_URL
77+
value: $(NODE_IP):8353
78+
- name: JAX_PLATFORMS
79+
value: tpu,cpu
80+
- name: ENABLE_PJRT_COMPATIBILITY
81+
value: "true"
82+
ports:
83+
- containerPort: 8081
84+
name: mxla
85+
nodeSelector:
86+
cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
87+
cloud.google.com/gke-tpu-topology: 4x4
88+
# [END gke_ai_ml_gke_ray_tpu_ray_job_tpu_v6e_multihost]

‎ai-ml/gke-ray/tpu/ray-job.tpu-v6e-singlehost.yaml‎

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,74 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
# [START gke_ai_ml_gke_ray_tpu_ray_job_tpu_v6e_singlehost]
16+
apiVersion: ray.io/v1
17+
kind: RayJob
18+
metadata:
19+
name: v6e-4-job
20+
spec:
21+
entrypoint: python ai-ml/gke-ray/tpu/tpu_list_devices.py
22+
runtimeEnvYAML: |
23+
working_dir: "https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/archive/refs/heads/main.zip"
24+
pip:
25+
- jax[tpu]==0.4.33
26+
- -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
27+
rayClusterSpec:
28+
rayVersion: '2.43.0'
29+
headGroupSpec:
30+
rayStartParams: {}
31+
template:
32+
spec:
33+
containers:
34+
- name: ray-head
35+
image: rayproject/ray:2.43.0-py310
36+
ports:
37+
- containerPort: 6379
38+
name: gcs-server
39+
- containerPort: 8265
40+
name: dashboard
41+
- containerPort: 10001
42+
name: client
43+
resources:
44+
limits:
45+
cpu: "8"
46+
memory: 40G
47+
requests:
48+
cpu: "8"
49+
memory: 40G
50+
workerGroupSpecs:
51+
- replicas: 1
52+
minReplicas: 1
53+
maxReplicas: 1
54+
numOfHosts: 1
55+
groupName: tpu-group
56+
rayStartParams: {}
57+
template:
58+
spec:
59+
containers:
60+
- name: ray-worker
61+
image: rayproject/ray:2.43.0-py310
62+
resources:
63+
limits:
64+
cpu: "24"
65+
google.com/tpu: "4"
66+
memory: 200G
67+
requests:
68+
cpu: "24"
69+
google.com/tpu: "4"
70+
memory: 200G
71+
nodeSelector:
72+
cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
73+
cloud.google.com/gke-tpu-topology: 2x2
74+
# [END gke_ai_ml_gke_ray_tpu_ray_job_tpu_v6e_singlehost]

‎ai-ml/gke-ray/tpu/tpu_list_devices.py‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import os
16+
import ray
17+
import jax
18+
19+
from jax.experimental import multihost_utils
20+
21+
ray.init()
22+
23+
@ray.remote(resources={"TPU": 4})
24+
def tpu_cores():
25+
multihost_utils.sync_global_devices("sync")
26+
print(f"TPU Worker: {os.environ.get('TPU_WORKER_ID')}")
27+
return f"TPU cores: {jax.device_count()}"
28+
29+
num_workers = int(ray.available_resources()["TPU"]) // 4
30+
print(f"Number of TPU Workers: {num_workers}")
31+
result = [tpu_cores.remote() for _ in range(num_workers)]
32+
print(ray.get(result))

0 commit comments

Comments
 (0)