[operator] fix TPU multi-host RayJob and RayCluster samples (ray-project#3733)

TPU v6e multihost requires at least `jax[tpu]==0.4.33`. This version doesn't
support py39 which is the version used in `rayproject/ray:2.46.0`. So we use
`rayproject/ray:2.46.0-py310` instead.

Signed-off-by: David Xia <david@davidxia.com>

Author: davidxia

ray-operator/config/samples/ray-cluster.tpu-v4-multihost.yaml

@@ -13,7 +13,7 @@ spec:
       spec:
         containers:
         - name: ray-head
-          image: rayproject/ray:2.46.0
+          image: rayproject/ray:2.46.0-py310
           imagePullPolicy: IfNotPresent
           resources:
             limits:
@@ -57,7 +57,7 @@ spec:
       spec:
         containers:
         - name: ray-worker
-          image: rayproject/ray:2.46.0
+          image: rayproject/ray:2.46.0-py310
           imagePullPolicy: IfNotPresent
           resources:
             limits:

ray-operator/config/samples/ray-cluster.tpu-v6e-16-multihost.yaml

@@ -10,7 +10,7 @@ spec:
       spec:
         containers:
         - name: ray-head
-          image: rayproject/ray:2.46.0
+          image: rayproject/ray:2.46.0-py310
           imagePullPolicy: IfNotPresent
           resources:
             limits:
@@ -41,7 +41,7 @@ spec:
       spec:
         containers:
         - name: ray-worker
-          image: rayproject/ray:2.46.0
+          image: rayproject/ray:2.46.0-py310
           imagePullPolicy: IfNotPresent
           resources:
             limits:

ray-operator/config/samples/ray-cluster.tpu-v6e-256-multihost.yaml

@@ -10,7 +10,7 @@ spec:
       spec:
         containers:
         - name: ray-head
-          image: rayproject/ray:2.46.0
+          image: rayproject/ray:2.46.0-py310
           imagePullPolicy: IfNotPresent
           resources:
             limits:
@@ -41,7 +41,7 @@ spec:
       spec:
         containers:
         - name: ray-worker
-          image: rayproject/ray:2.46.0
+          image: rayproject/ray:2.46.0-py310
           imagePullPolicy: IfNotPresent
           resources:
             limits:

ray-operator/config/samples/ray-job.tpu-v6e-16-multihost.yaml

@@ -7,7 +7,7 @@ spec:
   runtimeEnvYAML: |
     working_dir: "https://github.com/ray-project/kuberay/archive/master.zip"
     pip:
-      - jax[tpu]==0.4.33
+      - jax[tpu]==0.6.1
       - -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
   rayClusterSpec:
     rayVersion: '2.46.0'
@@ -17,7 +17,7 @@ spec:
         spec:
           containers:
           - name: ray-head
-            image: rayproject/ray:2.46.0
+            image: rayproject/ray:2.46.0-py310
             ports:
             - containerPort: 6379
               name: gcs-server
@@ -47,7 +47,7 @@ spec:
             runAsUser: 0
           containers:
           - name: ray-worker
-            image: rayproject/ray:2.46.0
+            image: rayproject/ray:2.46.0-py310
             resources:
               limits:
                 cpu: "24"

ray-operator/config/samples/ray-job.tpu-v6e-256-multihost.yaml

@@ -17,7 +17,7 @@ spec:
         spec:
           containers:
           - name: ray-head
-            image: rayproject/ray:2.46.0
+            image: rayproject/ray:2.46.0-py310
             ports:
             - containerPort: 6379
               name: gcs-server
@@ -46,7 +46,7 @@ spec:
         spec:
           containers:
           - name: ray-worker
-            image: rayproject/ray:2.46.0
+            image: rayproject/ray:2.46.0-py310
             resources:
               limits:
                 cpu: "24"

ray-operator/config/samples/ray-job.tpu-v6e-singlehost.yaml

@@ -7,7 +7,7 @@ spec:
   runtimeEnvYAML: |
     working_dir: "https://github.com/ray-project/kuberay/archive/master.zip"
     pip:
-      - jax[tpu]==0.4.33
+      - jax[tpu]==0.6.1
       - -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
   rayClusterSpec:
     rayVersion: '2.46.0'
@@ -17,7 +17,7 @@ spec:
         spec:
           containers:
           - name: ray-head
-            image: rayproject/ray:2.46.0
+            image: rayproject/ray:2.46.0-py310
             ports:
             - containerPort: 6379
               name: gcs-server
@@ -45,7 +45,7 @@ spec:
             runAsUser: 0
           containers:
           - name: ray-worker
-            image: rayproject/ray:2.46.0
+            image: rayproject/ray:2.46.0-py310
             resources:
               limits:
                 cpu: "24"

ray-operator/config/samples/tpu/tpu_list_devices.py

@@ -1,7 +1,6 @@
 import os
 import ray
 import jax
-import time
 
 from jax.experimental import multihost_utils
 
@@ -10,9 +9,8 @@
 @ray.remote(resources={"TPU": 4})
 def tpu_cores():
     multihost_utils.sync_global_devices("sync")
-    cores = "TPU cores:" + str(jax.device_count())
-    print("TPU Worker: " + os.environ.get("TPU_WORKER_ID"))
-    return cores
+    print(f"TPU Worker: {os.environ.get('TPU_WORKER_ID')}")
+    return f"TPU cores: {jax.device_count()}"
 
 num_workers = int(ray.available_resources()["TPU"]) // 4
 print(f"Number of TPU Workers: {num_workers}")