[Core] Make multi-node job fail fast when one fails, and output segment fault

docs/source/running-jobs/environment-variables.rst

@@ -94,6 +94,8 @@ and ``run`` stages have different environment variables available.
 Setup Stage
 ~~~~~~~~~~~
 
+Since setup commands always run on all nodes of a cluster, SkyPilot ensures both of these environment variables (the ranks and the IP list) never change across multiple setups on the same cluster.
+
 .. list-table::
   :widths: 20 70 10
   :header-rows: 1

sky/backends/cloud_vm_ray_backend.py

@@ -17,7 +17,6 @@
 import time
 import typing
 from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union
-import uuid
 
 import colorama
 import filelock
@@ -3129,19 +3128,19 @@ def _setup(self, handle: CloudVmRayResourceHandle, task: task_lib.Task,
         # Disable connection sharing for setup script to avoid old
         # connections being reused, which may cause stale ssh agent
         # forwarding.
-        ssh_credentials.pop('ssh_control_name')
-        runners = command_runner.SSHCommandRunner.make_runner_list(
-            ip_list, port_list=port_list, **ssh_credentials)
+        ssh_credentials.pop('ssh_control_name', None)
 
-        random_suffix = uuid.uuid4().hex[:8]
-        remote_setup_file_name = f'/tmp/sky_setup_{random_suffix}'
+        remote_setup_file_name = f'/tmp/sky_setup_{self.run_timestamp}'
         # Need this `-i` option to make sure `source ~/.bashrc` work
         setup_cmd = f'/bin/bash -i {remote_setup_file_name} 2>&1'
 
-        def _setup_node(runner: command_runner.SSHCommandRunner) -> None:
+        def _setup_node(node_id: int) -> None:
             setup_envs = task.envs.copy()
             setup_envs['SKYPILOT_SETUP_NODE_IPS'] = '\n'.join(internal_ips)
-            setup_envs['SKYPILOT_SETUP_NODE_RANK'] = str(runners.index(runner))
+            setup_envs['SKYPILOT_SETUP_NODE_RANK'] = str(node_id)
+            runner = command_runner.SSHCommandRunner(ip_list[node_id],
+                                                     port=port_list[node_id],
+                                                     **ssh_credentials)
             setup_script = log_lib.make_task_bash_script(setup,
                                                          env_vars=setup_envs)
             with tempfile.NamedTemporaryFile('w', prefix='sky_setup_') as f:
@@ -3198,7 +3197,7 @@ def error_message() -> str:
         # which can cause the program waiting for all the threads to finish,
         # even if some of them raise exceptions. We should replace it with
         # multi-process.
-        subprocess_utils.run_in_parallel(_setup_node, runners)
+        subprocess_utils.run_in_parallel(_setup_node, range(num_nodes))
 
         if detach_setup:
             # Only set this when setup needs to be run outside the self._setup()

sky/utils/subprocess_utils.py

@@ -4,7 +4,7 @@
 import random
 import subprocess
 import time
-from typing import Any, Callable, List, Optional, Tuple, Union
+from typing import Any, Callable, Iterable, List, Optional, Tuple, Union
 
 import colorama
 import psutil
@@ -50,7 +50,7 @@ def get_parallel_threads() -> int:
     return max(4, cpu_count - 1)
 
 
-def run_in_parallel(func: Callable, args: List[Any]) -> List[Any]:
+def run_in_parallel(func: Callable, args: Iterable[Any]) -> List[Any]:
     """Run a function in parallel on a list of arguments.
 
     The function should raise a CommandError if the command fails.

tests/test_smoke.py

@@ -1514,6 +1514,8 @@ def test_multi_node_failure(generic_cloud: str):
             # TODO(zhwu): we use multi-thread to run the commands in setup
             # commands in parallel, which makes it impossible to fail fast
             # when one of the nodes fails. We should fix this in the future.
+            # The --detach-setup version can fail fast, as the setup is
+            # submitted to the remote machine, which does not use multi-thread.
             # Refer to the comment in `subprocess_utils.run_in_parallel`.
             # f'sky launch -y -c {name} --cloud {generic_cloud} tests/test_yamls/failed_worker_setup.yaml && exit 1',  # Ensure the job setup failed.
             f'sky launch -y -c {name} --cloud {generic_cloud} --detach-setup tests/test_yamls/failed_worker_setup.yaml',