Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix v6e TPU Scripts and RayJob CRs #2447

Merged
merged 4 commits into from
Oct 16, 2024
Merged

Fix v6e TPU Scripts and RayJob CRs #2447

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
668889e
Fix tpu_list_device script and RayJob CRs
ryanaoleary Oct 16, 2024
5456ef6
Remove unneeded lifecycle spec
ryanaoleary Oct 16, 2024
c4a233c
Fix num_workers calculation
ryanaoleary Oct 16, 2024
1572b02
Fix entrypoint path
ryanaoleary Oct 16, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 1 addition & 7 deletions ray-operator/config/samples/ray-job.tpu-v6e-multihost.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,16 +1,14 @@
apiVersion: ray.io/v1
kind: RayJob
metadata:
name: ray-multi-host-job
name: v6e-256-job
spec:
entrypoint: python ray-operator/config/samples/ray-tpu-scripts/tpu_list_devices.py
runtimeEnvYAML: |
working_dir: "https://github.com/ray-project/kuberay/archive/master.zip"
pip:
- jax[tpu]==0.4.33
- -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
env_vars:
TPU_HOSTS: "64"
rayClusterSpec:
rayVersion: '2.37.0'
headGroupSpec:
Expand Down Expand Up @@ -50,10 +48,6 @@ spec:
containers:
- name: ray-worker
image: rayproject/ray:2.37.0-py310
lifecycle:
preStop:
exec:
command: [ "/bin/sh","-c","ray stop" ]
resources:
limits:
cpu: "100"
Expand Down
8 changes: 1 addition & 7 deletions ray-operator/config/samples/ray-job.tpu-v6e-singlehost.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,16 +1,14 @@
apiVersion: ray.io/v1
kind: RayJob
metadata:
name: ray-single-host-job
name: v6e-4-job
spec:
entrypoint: python ray-operator/config/samples/ray-tpu-scripts/tpu_list_devices.py
runtimeEnvYAML: |
working_dir: "https://github.com/ray-project/kuberay/archive/master.zip"
pip:
- jax[tpu]==0.4.33
- -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
env_vars:
TPU_HOSTS: "1"
rayClusterSpec:
rayVersion: '2.37.0'
headGroupSpec:
Expand Down Expand Up @@ -50,10 +48,6 @@ spec:
containers:
- name: ray-worker
image: rayproject/ray:2.37.0-py310
lifecycle:
preStop:
exec:
command: [ "/bin/sh","-c","ray stop" ]
resources:
limits:
cpu: "100"
Expand Down
13 changes: 3 additions & 10 deletions ray-operator/config/samples/tpu/tpu_list_devices.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,19 +1,12 @@
import ray
import jax

ray.init(
runtime_env={
"pip": [
"jax[tpu]==0.4.33",
"-f https://storage.googleapis.com/jax-releases/libtpu_releases.html",
]
}
)
ray.init()

@ray.remote(resources={"TPU": 4})
def tpu_cores():
import jax
return "TPU cores:" + str(jax.device_count())

num_workers = int(os.environ['TPU_HOSTS']) # Set in env of RayJob or RayCluster.
num_workers = int(ray.available_resources()["TPU"])
andrewsykim marked this conversation as resolved.
Show resolved Hide resolved
result = [tpu_cores.remote() for _ in range(num_workers)]
print(ray.get(result))
Loading