fix tpu bug (#2350)

* fix tpu bug

* format

* update

sky/backends/cloud_vm_ray_backend.py

@@ -2248,7 +2248,7 @@ def _update_stable_ssh_ports(self, max_attempts: int = 1) -> None:
             ports = [head_port] + worker_ports
         else:
             # Use port 22 for other clouds
-            ports = [22] * self.launched_nodes
+            ports = [22] * self.num_node_ips
         self.stable_ssh_ports = ports
 
     def _update_stable_cluster_ips(self, max_attempts: int = 1) -> None:
@@ -2350,6 +2350,16 @@ def head_ssh_port(self):
             return external_ssh_ports[0]
         return None
 
+    @property
+    def num_node_ips(self) -> int:
+        """Returns number of IPs of the cluster, correctly handling TPU Pod."""
+        is_tpu_vm_pod = tpu_utils.is_tpu_vm_pod(self.launched_resources)
+        if is_tpu_vm_pod:
+            num_ips = tpu_utils.get_num_tpu_devices(self.launched_resources)
+        else:
+            num_ips = self.launched_nodes
+        return num_ips
+
     def __setstate__(self, state):
         self._version = self._VERSION
 

sky/utils/tpu_utils.py

@@ -30,10 +30,9 @@ def is_tpu_vm_pod(resources: Optional[resources_lib.Resources]) -> bool:
     return acc not in ['tpu-v2-8', 'tpu-v3-8', 'tpu-v4-8']
 
 
-def get_num_tpu_devices(
-        resources: Optional[resources_lib.Resources]) -> Optional[int]:
+def get_num_tpu_devices(resources: Optional[resources_lib.Resources]) -> int:
     if resources is None or not is_tpu(resources):
-        return None
+        raise ValueError('resources must be a valid TPU resource.')
     acc, _ = list(resources.accelerators.items())[0]
     num_tpu_devices = int(int(acc.split('-')[2]) / 8)
     return num_tpu_devices