multithreaded ssh setup

sky/provision/kubernetes/instance.py

@@ -18,6 +18,7 @@
 from sky.utils import command_runner
 from sky.utils import common_utils
 from sky.utils import kubernetes_enums
+from sky.utils import subprocess_utils
 from sky.utils import ux_utils
 
 POLL_INTERVAL = 2
@@ -398,8 +399,7 @@ def _setup_ssh_in_pods(namespace: str, context: Optional[str],
         # See https://www.educative.io/answers/error-mesg-ttyname-failed-inappropriate-ioctl-for-device  # pylint: disable=line-too-long
         '$(prefix_cmd) sed -i "s/mesg n/tty -s \\&\\& mesg n/" ~/.profile;')
 
-    # TODO(romilb): Parallelize the setup of SSH in pods for multi-node clusters
-    for new_node in new_nodes:
+    def _setup_ssh_thread(new_node):
         pod_name = new_node.metadata.name
         runner = command_runner.KubernetesCommandRunner(
             ((namespace, context), pod_name))
@@ -411,6 +411,8 @@ def _setup_ssh_in_pods(namespace: str, context: Optional[str],
                                      stdout)
         logger.info(f'{"-"*20}End: Set up SSH in pod {pod_name!r} {"-"*20}')
 
+    subprocess_utils.run_in_parallel(_setup_ssh_thread, new_nodes)
+
 
 def _label_pod(namespace: str, context: Optional[str], pod_name: str,
                label: Dict[str, str]) -> None:

sky/provision/provisioner.py

@@ -28,6 +28,7 @@
 from sky.utils import common_utils
 from sky.utils import resources_utils
 from sky.utils import rich_utils
+from sky.utils import subprocess_utils
 from sky.utils import ux_utils
 
 # Do not use __name__ as we do not want to propagate logs to sky.provision,
@@ -365,21 +366,28 @@ def wait_for_ssh(cluster_info: provision_common.ClusterInfo,
     # use a queue for SSH querying
     ips = collections.deque(ip_list)
     ssh_ports = collections.deque(port_list)
-    while ips:
-        ip = ips.popleft()
-        ssh_port = ssh_ports.popleft()
-        success, stderr = waiter(ip, ssh_port, **ssh_credentials)
-        if not success:
-            ips.append(ip)
-            ssh_ports.append(ssh_port)
-            if time.time() - start > timeout:
+
+    def _retry_ssh_thread(ip_ssh_port: Tuple[str, int]):
+        ip, ssh_port = ip_ssh_port
+        success = False
+        while not success:
+            success, stderr = waiter(ip, ssh_port, **ssh_credentials)
+            if not success and time.time() - start > timeout:
                 with ux_utils.print_exception_no_traceback():
                     raise RuntimeError(
                         f'Failed to SSH to {ip} after timeout {timeout}s, with '
                         f'{stderr}')
             logger.debug('Retrying in 1 second...')
             time.sleep(1)
 
+    # try one node and multiprocess the rest
+    if ips:
+        ip = ips.popleft()
+        ssh_port = ssh_ports.popleft()
+        _retry_ssh_thread((ip, ssh_port))
+    subprocess_utils.run_in_parallel(_retry_ssh_thread,
+                                     list(zip(ips, ssh_ports)))
+
 
 def _post_provision_setup(
         cloud_name: str, cluster_name: resources_utils.ClusterName,