Repository: skypilot-org/skypilot
Pull request: [Core] Add support for detach setup
Closes #730.

This PR adds the following option to `sky.launch` and `sky launch`: `--async-setup`

Note:
1. The setup will be run on all the nodes in the clusters no matter how many nodes are required by the task so that all the nodes will be in the same state to avoid inconsistent behavior on different nodes after many setups are run. (There is still a barrier after setup, i.e. the setups should finish on all the nodes before the `run` is executed).
2. Added a JobStatus called `SETUP` and a life cycle of a job will be `INIT -> SETUP -> PENDING -> RUNNING -> SUCCEEDED/FAILED`
3. Discussion: Should we combine the `--async-setup` and `--detach-run` together to `--detach`, i.e. `setup` and `run` will detach together?
  Pro: simpler API for the user
  Con: Less control, some user's logic will need to be changed if they rely on the previous behavior.
  My thought: We can add the `--async-setup` first in this PR, to see how people feel about this, and defer the combination.

Example:
```
> sky launch -c async-setup4 -s --num-nodes 4 examples/minimal.yaml
Task from YAML spec: examples/minimal.yaml
I 11-05 15:21:13 optimizer.py:605] == Optimizer ==
I 11-05 15:21:13 optimizer.py:628] Estimated cost: $1.5 / hour
I 11-05 15:21:13 optimizer.py:628] 
I 11-05 15:21:13 optimizer.py:684] Considered resources (4 nodes):
I 11-05 15:21:13 optimizer.py:713] ------------------------------------------------------------------
I 11-05 15:21:13 optimizer.py:713]  CLOUD   INSTANCE      vCPUs   ACCELERATORS   COST ($)   CHOSEN   
I 11-05 15:21:13 optimizer.py:713] ------------------------------------------------------------------
I 11-05 15:21:13 optimizer.py:713]  AWS     m6i.2xlarge   8       -              1.54          ✔     
I 11-05 15:21:13 optimizer.py:713] ------------------------------------------------------------------
I 11-05 15:21:13 optimizer.py:713] 
Launching a new cluster 'async-setup4'. Proceed? [Y/n]: 
I 11-05 15:21:14 cloud_vm_ray_backend.py:2813] Creating a new cluster: "async-setup4" [4x AWS(m6i.2xlarge)].
I 11-05 15:21:14 cloud_vm_ray_backend.py:2813] Tip: to reuse an existing cluster, specify --cluster (-c). Run `sky status` to see existing clusters.
I 11-05 15:21:14 cloud_vm_ray_backend.py:958] To view detailed progress: tail -n100 -f /Users/zhwu/sky_logs/sky-2022-11-05-15-21-13-272410/provision.log
I 11-05 15:21:15 cloud_vm_ray_backend.py:1218] Launching on AWS us-west-2 (us-west-2d,us-west-2b,us-west-2a,us-west-2c)
I 11-05 15:22:45 log_utils.py:45] Head node is up.
I 11-05 15:23:55 cloud_vm_ray_backend.py:1326] Successfully provisioned or found existing head VM. Waiting for workers.
I 11-05 15:26:18 cloud_vm_ray_backend.py:1053] Successfully provisioned or found existing VMs.
I 11-05 15:26:38 cloud_vm_ray_backend.py:2079] Preparing setup for 4 nodes.
I 11-05 15:26:45 cloud_vm_ray_backend.py:2155] Job submitted with Job ID: 1
I 11-05 22:26:45 log_lib.py:388] Start streaming logs for job 1.
INFO: Tip: use Ctrl-C to exit log streaming (task will not be killed).
(setup pid=23992) running setup
(setup pid=22961, ip=172.31.37.209) running setup
(setup pid=22958, ip=172.31.43.180) running setup
(setup pid=22955, ip=172.31.37.178) running setup
INFO: Setup finished.
INFO: Waiting for task resources on 4 nodes. This will block if the cluster is full.
INFO: All task resources reserved.
INFO: Reserved IPs: ['172.31.43.180', '172.31.37.209', '172.31.45.100', '172.31.37.178']
(node-2 pid=23992) # conda environments:
(node-2 pid=23992) #
(node-2 pid=23992) base                  *  /opt/conda
(node-2 pid=23992) pytorch                  /opt/conda/envs/pytorch
(node-0 pid=22958, ip=172.31.43.180) # conda environments:
(node-0 pid=22958, ip=172.31.43.180) #
(node-0 pid=22958, ip=172.31.43.180) base                  *  /opt/conda
(node-0 pid=22958, ip=172.31.43.180) pytorch                  /opt/conda/envs/pytorch
(node-3 pid=22955, ip=172.31.37.178) # conda environments:
(node-3 pid=22955, ip=172.31.37.178) #
(node-3 pid=22955, ip=172.31.37.178) base                  *  /opt/conda
(node-3 pid=22955, ip=172.31.37.178) pytorch                  /opt/conda/envs/pytorch
(node-1 pid=22961, ip=172.31.37.209) # conda environments:
(node-1 pid=22961, ip=172.31.37.209) #
(node-1 pid=22961, ip=172.31.37.209) base                  *  /opt/conda
(node-1 pid=22961, ip=172.31.37.209) pytorch                  /opt/conda/envs/pytorch
(node-1 pid=22961, ip=172.31.37.209) 
(node-0 pid=22958, ip=172.31.43.180) 
(node-2 pid=23992) 
(node-3 pid=22955, ip=172.31.37.178) 
INFO: Job finished (status: SUCCEEDED).
Shared connection to 35.88.43.10 closed.
I 11-05 15:28:32 cloud_vm_ray_backend.py:2184] Job ID: 1
I 11-05 15:28:32 cloud_vm_ray_backend.py:2184] To cancel the job:       sky cancel async-setup4 1
I 11-05 15:28:32 cloud_vm_ray_backend.py:2184] To stream the logs:      sky logs async-setup4 1
I 11-05 15:28:32 cloud_vm_ray_backend.py:2184] To view the job queue:   sky queue async-setup4
I 11-05 15:28:32 cloud_vm_ray_backend.py:2297] 
I 11-05 15:28:32 cloud_vm_ray_backend.py:2297] Cluster name: async-setup4
I 11-05 15:28:32 cloud_vm_ray_backend.py:2297] To log into the head VM: ssh async-setup4
I 11-05 15:28:32 cloud_vm_ray_backend.py:2297] To submit a job:         sky exec async-setup4 yaml_file
I 11-05 15:28:32 cloud_vm_ray_backend.py:2297] To stop the cluster:     sky stop async-setup4
I 11-05 15:28:32 cloud_vm_ray_backend.py:2297] To teardown the cluster: sky down async-setup4
Clusters
NAME          LAUNCHED    RESOURCES                        STATUS   AUTOSTOP  COMMAND                        
async-setup4  2 mins ago  4x AWS(m6i.2xlarge)              UP       -         sky launch -c async-setup4...  
smoke-test    2 days ago  1x AWS(m6i.xlarge)               STOPPED  -         sky start smoke-test           
workspace     1 week ago  1x AWS(m6i.large, disk_size=50)  UP       -         sky launch -c workspace -...   

Managed spot controller (will be autostopped if idle for 10min)
NAME                          LAUNCHED    RESOURCES                          STATUS   AUTOSTOP  COMMAND                  
sky-spot-controller-9ce1ce58  5 days ago  1x AWS(m6i.2xlarge, disk_size=50)  STOPPED  10m       sky spot launch echo hi  

Local clusters:
NAME         USER  HEAD_IP  RESOURCES  COMMAND  
gpu-cluster  -     -        -          -        
```

TODO:
- [x] Add tests for this PR.

Tested:
- [x] `sky launch -c async-setup -s examples/minimal.yaml` and the status shows `SETUP` correctly during setup and transit to `RUNNING` later.
- [x] `sky launch -c async-setup4 -s --num-nodes 4 examples/minimal.yaml`
- [x] `./tests/run_smoke_tests.sh test_n_node_job_queue`
- [x] `./tests/run_smoke_tests.sh`
- [x] `sky launch -c async-setup -s  examples/minimal.yaml` (with `exit 1` in the setup`)
```
I 11-07 00:01:25 log_utils.py:45] Head node is up.
I 11-07 00:02:46 cloud_vm_ray_backend.py:1063] Successfully provisioned or found existing VM.
I 11-07 00:02:47 cloud_vm_ray_backend.py:2094] Preparing setup for 1 node.
I 11-07 00:02:50 cloud_vm_ray_backend.py:2173] Job submitted with Job ID: 1
I 11-07 08:02:51 log_lib.py:388] Start streaming logs for job 1.
INFO: Tip: use Ctrl-C to exit log streaming (task will not be killed).
(setup pid=24911) running setup
ERROR: Job 1 setup failed with return code list: [1]
INFO: Job finished (status: FAILED_SETUP).
Shared connection to 35.86.185.118 closed.
```
```
> sky queue async-setup
Fetching and parsing job queue...

Job queue of cluster async-setup
ID  NAME     SUBMITTED  STARTED  DURATION  RESOURCES     STATUS        LOG                                        
1   minimal  1 min ago  -        -         1x [CPU:0.5]  FAILED_SETUP  ~/sky_logs/sky-2022-11-06-23-59-54-232565  
```

Problem statement:
Support asynchronous setup and file_mounts
Our current setup and file_mounts require the user's laptop to have an `ssh` connection with the cluster, which means the setup or file_mounts will be stopped if the user stop the laptop. When the setup takes a long time, it is easy for the user to fail the setup.

Workaround: place setup in the run section.

Hints from issue discussion:
This falls into the asynchronous story, right?
For the setup, I think it may be a little bit different. For the asynchronous story, we assume there is a process running in the background on the laptop connecting with the cluster and handling the setup/file_mounts.
However, the `setup` can be run on the cluster independently, like the `run`, without the existence of the user's laptop. That is to say, we may need to submit the setup script as a job, instead of an interactive process.
@dementrock ran into this again:
> This would be a big usability improvement for me to be able to run the setup process without having to guard my laptop to be open & have good internet connection since I have kind of a long-ish setup right now (1~2 hours).

Commit message: Add test for async setup

Patch to review:
diff --git a/examples/job_queue/job_multinode.yaml b/examples/job_queue/job_multinode.yaml
index ad016154e1d..ff291589001 100644
--- a/examples/job_queue/job_multinode.yaml
+++ b/examples/job_queue/job_multinode.yaml
@@ -20,7 +20,7 @@ setup: |
 run: |
   timestamp=$(date +%s)
   conda env list
-  for i in {1..120}; do
+  for i in {1..240}; do
     echo "$timestamp $i"
     sleep 1
   done
diff --git a/examples/job_queue/job_multinode_long_setup.yaml b/examples/job_queue/job_multinode_long_setup.yaml
new file mode 100644
index 00000000000..44f30105338
--- /dev/null
+++ b/examples/job_queue/job_multinode_long_setup.yaml
@@ -0,0 +1,27 @@
+# A task runs on an existing multinode cluster.
+#
+# Runs a task that requires multinode.
+#
+# Usage:
+#   sky launch -c mjq cluster_multinode.yaml
+#   sky exec mjq job_multinode.yaml
+#   sky exec mjq job.yaml
+
+name: job_multinode
+
+resources:
+  accelerators: K80:0.5
+
+num_nodes: 2
+
+setup: |
+  echo "running setup"
+  sleep 80
+
+run: |
+  timestamp=$(date +%s)
+  conda env list
+  for i in {1..120}; do
+    echo "$timestamp $i"
+    sleep 1
+  done
diff --git a/sky/backends/backend.py b/sky/backends/backend.py
index 0a6d816f291..2b87d7a82f6 100644
--- a/sky/backends/backend.py
+++ b/sky/backends/backend.py
@@ -66,20 +66,24 @@ def sync_file_mounts(
 
     @timeline.event
     @usage_lib.messages.usage.update_runtime('setup')
-    def setup(self, handle: ResourceHandle, task: 'task_lib.Task') -> None:
-        return self._setup(handle, task)
+    def setup(self, handle: ResourceHandle, task: 'task_lib.Task',
+              async_setup: bool) -> None:
+        return self._setup(handle, task, async_setup)
 
     def add_storage_objects(self, task: 'task_lib.Task') -> None:
         raise NotImplementedError
 
     @timeline.event
     @usage_lib.messages.usage.update_runtime('execute')
-    def execute(self, handle: ResourceHandle, task: 'task_lib.Task',
-                detach_run: bool) -> None:
+    def execute(self,
+                handle: ResourceHandle,
+                task: 'task_lib.Task',
+                detach_run: bool,
+                setup_cmd: Optional[str] = None) -> None:
         usage_lib.record_cluster_name_for_current_operation(
             handle.get_cluster_name())
         usage_lib.messages.usage.update_actual_task(task)
-        return self._execute(handle, task, detach_run)
+        return self._execute(handle, task, detach_run, setup_cmd=setup_cmd)
 
     @timeline.event
     def post_execute(self, handle: ResourceHandle, down: bool) -> None:
@@ -123,11 +127,15 @@ def _sync_file_mounts(
     ) -> None:
         raise NotImplementedError
 
-    def _setup(self, handle: ResourceHandle, task: 'task_lib.Task') -> None:
+    def _setup(self, handle: ResourceHandle, task: 'task_lib.Task',
+               async_setup: bool) -> None:
         raise NotImplementedError
 
-    def _execute(self, handle: ResourceHandle, task: 'task_lib.Task',
-                 detach_run: bool) -> None:
+    def _execute(self,
+                 handle: ResourceHandle,
+                 task: 'task_lib.Task',
+                 detach_run: bool,
+                 setup_cmd: Optional[str] = None) -> None:
         raise NotImplementedError
 
     def _post_execute(self, handle: ResourceHandle, down: bool) -> None:
diff --git a/sky/backends/cloud_vm_ray_backend.py b/sky/backends/cloud_vm_ray_backend.py
index 9c185b45749..be878c2071e 100644
--- a/sky/backends/cloud_vm_ray_backend.py
+++ b/sky/backends/cloud_vm_ray_backend.py
@@ -92,6 +92,9 @@
 
 _TPU_NOT_FOUND_ERROR = 'ERROR: (gcloud.compute.tpus.delete) NOT_FOUND'
 
+_CTRL_C_TIP_MESSAGE = ('INFO: Tip: use Ctrl-C to exit log streaming '
+                       '(task will not be killed).')
+
 _MAX_RAY_UP_RETRY = 5
 
 _JOB_ID_PATTERN = re.compile(r'Job ID: ([0-9]+)')
@@ -148,6 +151,9 @@ def __init__(self):
     def add_prologue(self,
                      job_id: int,
                      spot_task: Optional['task_lib.Task'] = None,
+                     setup_cmd: Optional[str] = None,
+                     envs: Optional[Dict[str, str]] = None,
+                     setup_log_path: Optional[str] = None,
                      is_local: bool = False) -> None:
         assert not self._has_prologue, 'add_prologue() called twice?'
         self._has_prologue = True
@@ -182,12 +188,11 @@ def add_prologue(self,
             from sky.utils import log_utils
 
             SKY_REMOTE_WORKDIR = {constants.SKY_REMOTE_WORKDIR!r}
-            job_lib.set_status({job_id!r}, job_lib.JobStatus.PENDING)
 
             ray.init(address={ray_address!r}, namespace='__sky__{job_id}__', log_to_driver=True)
-
             run_fn = None
-            futures = []"""),
+            futures = []
+            """),
             # FIXME: This is a hack to make sure that the functions can be found
             # by ray.remote. This should be removed once we have a better way to
             # specify dependencies for ray.
@@ -197,6 +202,36 @@ def add_prologue(self,
             inspect.getsource(log_lib.add_ray_env_vars),
             inspect.getsource(log_lib.run_bash_command_with_log),
             'run_bash_command_with_log = ray.remote(run_bash_command_with_log)',
+            f'setup_cmd = {setup_cmd!r}'
+        ]
+        if setup_cmd is not None:
+            self._code += [
+                textwrap.dedent(f"""\
+                _SETUP_CPUS = 0.0001
+                job_lib.set_status({job_id!r}, job_lib.JobStatus.SETUP)
+                print({_CTRL_C_TIP_MESSAGE!r}, file=sys.stderr, flush=True)
+                total_num_nodes = len(ray.nodes())
+                bundles = [{{"CPU": _SETUP_CPUS}} for _ in range(total_num_nodes)]
+                pg = ray.util.placement_group(bundles, strategy='STRICT_SPREAD')
+                ray.get(pg.ready())
+                setup_workers = [run_bash_command_with_log \\
+                    .options(name='setup', num_cpus=_SETUP_CPUS, placement_group=pg, placement_group_bundle_index=i) \\
+                    .remote(
+                        setup_cmd,
+                        os.path.expanduser({setup_log_path!r}),
+                        getpass.getuser(),
+                        job_id={self.job_id},
+                        env_vars={envs!r},
+                        stream_logs=True,
+                        with_ray=True,
+                        use_sudo={is_local},
+                    ) for i in range(total_num_nodes)]
+                ray.get(setup_workers)
+                time.sleep(1)
+                print('INFO: Setup finished.', file=sys.stderr, flush=True)""")
+            ]
+        self._code += [
+            f'job_lib.set_status({job_id!r}, job_lib.JobStatus.PENDING)',
         ]
         if spot_task is not None:
             # Add the spot job to spot queue table.
@@ -247,15 +282,15 @@ def add_gang_scheduling_placement_group(
                     **gpu_dict,
                 })
 
-        pack_mode = 'STRICT_SPREAD'
         self._code += [
             textwrap.dedent(f"""\
-                pg = ray_util.placement_group({json.dumps(bundles)}, {pack_mode!r})
+                pg = ray_util.placement_group({json.dumps(bundles)}, 'STRICT_SPREAD')
                 plural = 's' if {num_nodes} > 1 else ''
-                node_str = f'{num_nodes} node' + plural + '.'
-                print('INFO: Tip: use Ctrl-C to exit log streaming (task will not be killed).\\n'
-                      'INFO: Waiting for task resources on ' + node_str +
-                      ' This will block if the cluster is full.',
+                node_str = f'{num_nodes} node{{plural}}'
+                
+                message = '' if setup_cmd is not None else {_CTRL_C_TIP_MESSAGE!r} + '\\n'
+                message += f'INFO: Waiting for task resources on {{node_str}}. This will block if the cluster is full.'
+                print(message,
                       file=sys.stderr,
                       flush=True)
                 # FIXME: This will print the error message from autoscaler if
@@ -1961,7 +1996,8 @@ def _sync_file_mounts(
         self._execute_file_mounts(handle, all_file_mounts)
         self._execute_storage_mounts(handle, storage_mounts)
 
-    def _setup(self, handle: ResourceHandle, task: task_lib.Task) -> None:
+    def _setup(self, handle: ResourceHandle, task: task_lib.Task,
+               async_setup: bool) -> Optional[str]:
         start = time.time()
         style = colorama.Style
         fore = colorama.Fore
@@ -1988,17 +2024,20 @@ def _setup(self, handle: ResourceHandle, task: task_lib.Task) -> None:
             runners = command_runner.SSHCommandRunner.make_runner_list(
                 ip_list, *ssh_credentials)
 
+            # Need this `-i` option to make sure `source ~/.bashrc` work
+            setup_cmd = f'/bin/bash -i /tmp/{setup_file} 2>&1'
+
             def _setup_node(runner: command_runner.SSHCommandRunner) -> int:
                 runner.rsync(source=setup_sh_path,
                              target=f'/tmp/{setup_file}',
                              up=True,
                              stream_logs=False)
-                # Need this `-i` option to make sure `source ~/.bashrc` work
-                cmd = f'/bin/bash -i /tmp/{setup_file} 2>&1'
+                if async_setup:
+                    return
                 setup_log_path = os.path.join(self.log_dir,
                                               f'setup-{runner.ip}.log')
                 returncode = runner.run(
-                    cmd,
+                    setup_cmd,
                     log_path=setup_log_path,
                     process_stream=False,
                 )
@@ -2031,14 +2070,23 @@ def error_message() -> str:
                     return err_msg
 
                 subprocess_utils.handle_returncode(returncode=returncode,
-                                                   command=cmd,
+                                                   command=setup_cmd,
                                                    error_msg=error_message)
 
             num_nodes = len(ip_list)
             plural = 's' if num_nodes > 1 else ''
-            logger.info(f'{fore.CYAN}Running setup on {num_nodes} node{plural}.'
-                        f'{style.RESET_ALL}')
+            if async_setup:
+                logger.info(
+                    f'{fore.CYAN}Preparing setup for {num_nodes} node{plural}.'
+                    f'{style.RESET_ALL}')
+            else:
+                logger.info(
+                    f'{fore.CYAN}Running setup on {num_nodes} node{plural}.'
+                    f'{style.RESET_ALL}')
             subprocess_utils.run_in_parallel(_setup_node, runners)
+
+        if async_setup:
+            return setup_cmd
         logger.info(f'{fore.GREEN}Setup completed.{style.RESET_ALL}')
         end = time.time()
         logger.debug(f'Setup took {end - start} seconds.')
@@ -2212,6 +2260,7 @@ def _execute(
         handle: ResourceHandle,
         task: task_lib.Task,
         detach_run: bool,
+        setup_cmd: Optional[str] = None,
     ) -> None:
         if task.run is None:
             logger.info('Run commands not specified or empty.')
@@ -2226,10 +2275,12 @@ def _execute(
         is_tpu_vm_pod = tpu_utils.is_tpu_vm_pod(handle.launched_resources)
         # Case: task_lib.Task(run, num_nodes=N) or TPU VM Pods
         if task.num_nodes > 1 or is_tpu_vm_pod:
-            self._execute_task_n_nodes(handle, task, job_id, detach_run)
+            self._execute_task_n_nodes(handle, task, job_id, detach_run,
+                                       setup_cmd)
         else:
             # Case: task_lib.Task(run, num_nodes=1)
-            self._execute_task_one_node(handle, task, job_id, detach_run)
+            self._execute_task_one_node(handle, task, job_id, detach_run,
+                                        setup_cmd)
 
     def _post_execute(self, handle: ResourceHandle, down: bool) -> None:
         colorama.init()
@@ -2981,9 +3032,12 @@ def _execute_storage_mounts(self, handle: ResourceHandle,
         end = time.time()
         logger.debug(f'Storage mount sync took {end - start} seconds.')
 
-    def _execute_task_one_node(self, handle: ResourceHandle,
-                               task: task_lib.Task, job_id: int,
-                               detach_run: bool) -> None:
+    def _execute_task_one_node(self,
+                               handle: ResourceHandle,
+                               task: task_lib.Task,
+                               job_id: int,
+                               detach_run: bool,
+                               setup_cmd: Optional[str] = None) -> None:
         # Launch the command as a Ray task.
         log_dir = os.path.join(self.log_dir, 'tasks')
         log_path = os.path.join(log_dir, 'run.log')
@@ -2994,6 +3048,9 @@ def _execute_task_one_node(self, handle: ResourceHandle,
         is_local = isinstance(handle.launched_resources.cloud, clouds.Local)
         codegen.add_prologue(job_id,
                              spot_task=task.spot_task,
+                             setup_cmd=setup_cmd,
+                             envs=task.envs,
+                             setup_log_path=os.path.join(log_dir, 'setup.log'),
                              is_local=is_local)
         codegen.add_gang_scheduling_placement_group(1, accelerator_dict)
 
@@ -3020,8 +3077,12 @@ def _execute_task_one_node(self, handle: ResourceHandle,
                                 executable='python3',
                                 detach_run=detach_run)
 
-    def _execute_task_n_nodes(self, handle: ResourceHandle, task: task_lib.Task,
-                              job_id: int, detach_run: bool) -> None:
+    def _execute_task_n_nodes(self,
+                              handle: ResourceHandle,
+                              task: task_lib.Task,
+                              job_id: int,
+                              detach_run: bool,
+                              setup_cmd: Optional[str] = None) -> None:
         # Strategy:
         #   ray.init(...)
         #   for node:
@@ -3051,6 +3112,9 @@ def _execute_task_n_nodes(self, handle: ResourceHandle, task: task_lib.Task,
         is_local = isinstance(handle.launched_resources.cloud, clouds.Local)
         codegen.add_prologue(job_id,
                              spot_task=task.spot_task,
+                             setup_cmd=setup_cmd,
+                             envs=task.envs,
+                             setup_log_path=os.path.join(log_dir, 'setup.log'),
                              is_local=is_local)
         codegen.add_gang_scheduling_placement_group(
             num_actual_nodes,
diff --git a/sky/benchmark/benchmark_utils.py b/sky/benchmark/benchmark_utils.py
index 969fe1d1ee4..f2b71a62f48 100644
--- a/sky/benchmark/benchmark_utils.py
+++ b/sky/benchmark/benchmark_utils.py
@@ -315,7 +315,7 @@ def _update_benchmark_result(benchmark_result: Dict[str, Any]) -> Optional[str]:
 
     # Update the benchmark status.
     if (cluster_status == global_user_state.ClusterStatus.INIT or
-            job_status in [job_lib.JobStatus.INIT, job_lib.JobStatus.PENDING]):
+            job_status < job_lib.JobStatus.RUNNING):
         benchmark_status = benchmark_state.BenchmarkStatus.INIT
     elif job_status == job_lib.JobStatus.RUNNING:
         benchmark_status = benchmark_state.BenchmarkStatus.RUNNING
diff --git a/sky/cli.py b/sky/cli.py
index bbe46783c3a..79e5dab5554 100644
--- a/sky/cli.py
+++ b/sky/cli.py
@@ -666,6 +666,7 @@ def _launch_with_confirm(
     *,
     dryrun: bool,
     detach_run: bool,
+    async_setup: bool = False,
     no_confirm: bool = False,
     idle_minutes_to_autostop: Optional[int] = None,
     down: bool = False,  # pylint: disable=redefined-outer-name
@@ -726,6 +727,7 @@ def _launch_with_confirm(
             dryrun=dryrun,
             stream_logs=True,
             cluster_name=cluster,
+            async_setup=async_setup,
             detach_run=detach_run,
             backend=backend,
             idle_minutes_to_autostop=idle_minutes_to_autostop,
@@ -1057,6 +1059,11 @@ def cli():
               default=False,
               is_flag=True,
               help='If True, do not actually run the job.')
+@click.option('--async-setup',
+              '-s',
+              default=False,
+              is_flag=True,
+              help='if True, run setup asynchronously.')
 @click.option('--detach-run',
               '-d',
               default=False,
@@ -1116,7 +1123,6 @@ def cli():
               required=False,
               help='Skip confirmation prompt.')
 @click.option('--no-setup',
-              '-n',
               is_flag=True,
               default=False,
               required=False,
@@ -1126,6 +1132,7 @@ def launch(
     entrypoint: str,
     cluster: Optional[str],
     dryrun: bool,
+    async_setup: bool,
     detach_run: bool,
     backend_name: Optional[str],
     name: Optional[str],
@@ -1190,6 +1197,7 @@ def launch(
         backend,
         cluster,
         dryrun=dryrun,
+        async_setup=async_setup,
         detach_run=detach_run,
         no_confirm=yes,
         idle_minutes_to_autostop=idle_minutes_to_autostop,
diff --git a/sky/execution.py b/sky/execution.py
index 5cf12e6d960..7bc07ba6e5f 100644
--- a/sky/execution.py
+++ b/sky/execution.py
@@ -114,6 +114,7 @@ def _execute(
     optimize_target: OptimizeTarget = OptimizeTarget.COST,
     stages: Optional[List[Stage]] = None,
     cluster_name: Optional[str] = None,
+    async_setup: bool = False,
     detach_run: bool = False,
     idle_minutes_to_autostop: Optional[int] = None,
     no_setup: bool = False,
@@ -146,6 +147,7 @@ def _execute(
         skipping all setup steps.
       cluster_name: Name of the cluster to create/reuse.  If None,
         auto-generate a name.
+      async_setup: bool; whether to run the setup asynchronously.
       detach_run: bool; whether to detach the process after the job submitted.
       idle_minutes_to_autostop: int; if provided, the cluster will be set to
         autostop after this many minutes of idleness.
@@ -237,10 +239,11 @@ def _execute(
             backend.sync_file_mounts(handle, task.file_mounts,
                                      task.storage_mounts)
 
+        setup_cmd = None
         if no_setup:
             logger.info('Setup commands skipped.')
         elif Stage.SETUP in stages:
-            backend.setup(handle, task)
+            setup_cmd = backend.setup(handle, task, async_setup=async_setup)
 
         if Stage.PRE_EXEC in stages:
             if idle_minutes_to_autostop is not None:
@@ -251,7 +254,7 @@ def _execute(
         if Stage.EXEC in stages:
             try:
                 global_user_state.update_last_use(handle.get_cluster_name())
-                backend.execute(handle, task, detach_run)
+                backend.execute(handle, task, detach_run, setup_cmd=setup_cmd)
             finally:
                 # Enables post_execute() to be run after KeyboardInterrupt.
                 backend.post_execute(handle, down)
@@ -289,6 +292,7 @@ def launch(
     stream_logs: bool = True,
     backend: Optional[backends.Backend] = None,
     optimize_target: OptimizeTarget = OptimizeTarget.COST,
+    async_setup: bool = False,
     detach_run: bool = False,
     no_setup: bool = False,
 ) -> None:
@@ -328,7 +332,8 @@ def launch(
             (CloudVMRayBackend).
         optimize_target: target to optimize for. Choices: OptimizeTarget.COST,
             OptimizeTarget.TIME.
-        detach_run: If True, run setup first (blocking), then detach from the
+        async_setup: If True, run setup asynchronously.
+        detach_run: If True, run setup first, then detach from the
             job's execution.
         no_setup: if True, do not re-run setup commands.
 
@@ -355,6 +360,7 @@ def launch(
         retry_until_up=retry_until_up,
         optimize_target=optimize_target,
         cluster_name=cluster_name,
+        async_setup=async_setup,
         detach_run=detach_run,
         idle_minutes_to_autostop=idle_minutes_to_autostop,
         no_setup=no_setup,
diff --git a/sky/skylet/job_lib.py b/sky/skylet/job_lib.py
index 24514680613..19fa85699d7 100644
--- a/sky/skylet/job_lib.py
+++ b/sky/skylet/job_lib.py
@@ -78,6 +78,8 @@ class JobStatus(enum.Enum):
     # not started yet. skylet can transit the state from INIT to FAILED
     # directly, if the ray program fails to start.
     INIT = 'INIT'
+    # Running the user's setup script.
+    SETUP = 'SETUP'
     # The job is waiting for the required resources. (`ray job status`
     # shows RUNNING as the generated ray program has started, but blocked
     # by the placement constraints.)
@@ -94,7 +96,7 @@ class JobStatus(enum.Enum):
 
     @classmethod
     def nonterminal_statuses(cls) -> List['JobStatus']:
-        return [cls.INIT, cls.PENDING, cls.RUNNING]
+        return [cls.INIT, cls.SETUP, cls.PENDING, cls.RUNNING]
 
     def is_terminal(self):
         return self not in self.nonterminal_statuses()
@@ -115,7 +117,7 @@ def __lt__(self, other):
     # reserved resources and actually started running: it will set the
     # status in the DB to RUNNING.
     'PENDING': JobStatus.INIT,
-    'RUNNING': JobStatus.PENDING,
+    'RUNNING': JobStatus.SETUP,
     'SUCCEEDED': JobStatus.SUCCEEDED,
     'FAILED': JobStatus.FAILED,
     'STOPPED': JobStatus.CANCELLED,
@@ -390,7 +392,7 @@ def update_job_status(job_owner: str,
                 # already been set to later state by the job. We skip the
                 # update.
                 # 2. _RAY_TO_JOB_STATUS_MAP would map `ray job status`'s
-                # `RUNNING` to our JobStatus.PENDING; if a job has already been
+                # `RUNNING` to our JobStatus.SETYO; if a job has already been
                 # set to JobStatus.RUNNING by the generated ray program,
                 # `original_status` (job status from our DB) would already have
                 # that value. So we take the max here to keep it at RUNNING.
@@ -475,7 +477,7 @@ def dump_job_queue(username: Optional[str], all_jobs: bool) -> str:
         username: The username to show jobs for. Show all the users if None.
         all_jobs: Whether to show all jobs, not just the pending/running ones.
     """
-    status_list = [JobStatus.PENDING, JobStatus.RUNNING]
+    status_list = [JobStatus.SETUP, JobStatus.PENDING, JobStatus.RUNNING]
     if all_jobs:
         status_list = None
 
@@ -508,7 +510,8 @@ def cancel_jobs(job_owner: str, jobs: Optional[List[int]]) -> None:
     # Update the status of the jobs to avoid setting the status of stale
     # jobs to CANCELLED.
     if jobs is None:
-        job_records = _get_jobs(None, [JobStatus.PENDING, JobStatus.RUNNING])
+        job_records = _get_jobs(
+            None, [JobStatus.SETUP, JobStatus.PENDING, JobStatus.RUNNING])
     else:
         job_records = _get_jobs_by_ids(jobs)
 
@@ -528,7 +531,9 @@ def cancel_jobs(job_owner: str, jobs: Optional[List[int]]) -> None:
             logger.warning(str(e))
             continue
 
-        if job['status'] in [JobStatus.PENDING, JobStatus.RUNNING]:
+        if job['status'] in [
+                JobStatus.SETUP, JobStatus.PENDING, JobStatus.RUNNING
+        ]:
             set_status(job['job_id'], JobStatus.CANCELLED)
 
 
diff --git a/sky/skylet/log_lib.py b/sky/skylet/log_lib.py
index a8ede3f570c..b1e14994242 100644
--- a/sky/skylet/log_lib.py
+++ b/sky/skylet/log_lib.py
@@ -351,7 +351,8 @@ def _follow_job_logs(file,
             # Auto-exit the log tailing, if the job has finished. Check
             # the job status before query again to avoid unfinished logs.
             if status not in [
-                    job_lib.JobStatus.RUNNING, job_lib.JobStatus.PENDING
+                    job_lib.JobStatus.SETUP, job_lib.JobStatus.PENDING,
+                    job_lib.JobStatus.RUNNING
             ]:
                 if wait_last_logs:
                     # Wait all the logs are printed before exit.
@@ -398,11 +399,7 @@ def tail_logs(job_owner: str,
     # Wait for the log to be written. This is needed due to the `ray submit`
     # will take some time to start the job and write the log.
     retry_cnt = 0
-    while status in [
-            job_lib.JobStatus.INIT,
-            job_lib.JobStatus.PENDING,
-            job_lib.JobStatus.RUNNING,
-    ]:
+    while status is not None and not status.is_terminal():
         retry_cnt += 1
         if os.path.exists(log_path) and status != job_lib.JobStatus.INIT:
             break
@@ -419,7 +416,9 @@ def tail_logs(job_owner: str,
 
     start_stream_at = 'INFO: Tip: use Ctrl-C to exit log'
     if follow and status in [
-            job_lib.JobStatus.RUNNING, job_lib.JobStatus.PENDING
+            job_lib.JobStatus.SETUP,
+            job_lib.JobStatus.PENDING,
+            job_lib.JobStatus.RUNNING,
     ]:
         # Not using `ray job logs` because it will put progress bar in
         # multiple lines.
diff --git a/sky/spot/recovery_strategy.py b/sky/spot/recovery_strategy.py
index 6c33cc25a89..7cc439397df 100644
--- a/sky/spot/recovery_strategy.py
+++ b/sky/spot/recovery_strategy.py
@@ -222,8 +222,7 @@ def _launch(self, max_retry=3, raise_on_failure=True) -> Optional[float]:
                     continue
 
                 # Check the job status until it is not in initialized status
-                if status not in (None, job_lib.JobStatus.INIT,
-                                  job_lib.JobStatus.PENDING):
+                if status is not None and job_lib.JobStatus.PENDING < status:
                     try:
                         launch_time = spot_utils.get_job_timestamp(
                             self.backend, self.cluster_name, get_end_time=False)
diff --git a/tests/test_smoke.py b/tests/test_smoke.py
index 09611c77049..40a1aaaaf25 100644
--- a/tests/test_smoke.py
+++ b/tests/test_smoke.py
@@ -314,10 +314,12 @@ def test_n_node_job_queue():
             f'sky launch -y -c {name} examples/job_queue/cluster_multinode.yaml',
             f'sky exec {name} -n {name}-1 -d examples/job_queue/job_multinode.yaml',
             f'sky exec {name} -n {name}-2 -d examples/job_queue/job_multinode.yaml',
-            f'sky exec {name} -n {name}-3 -d examples/job_queue/job_multinode.yaml',
-            f'sky queue {name} | grep {name}-1 | grep RUNNING',
-            f'sky queue {name} | grep {name}-2 | grep RUNNING',
-            f'sky queue {name} | grep {name}-3 | grep PENDING',
+            f'sky launch -c {name} -n {name}-3 --async-setup -d examples/job_queue/job_multinode_long_setup.yaml',
+            f's=$(sky queue {name} | grep {name}-1) && echo $s && echo $s | grep RUNNING',
+            f's=$(sky queue {name} | grep {name}-2) && echo $s && echo $s | grep RUNNING',
+            f's=$(sky queue {name} | grep {name}-3) && echo $s && echo $s | grep SETUP',
+            'sleep 90',
+            f's=$(sky queue {name} | grep {name}-3) && echo $s && echo $s | grep PENDING',
             f'sky cancel {name} 1',
             'sleep 5',
             f'sky queue {name} | grep {name}-3 | grep RUNNING',


Code review comments:
diff_hunk: @@ -66,20 +66,24 @@ def sync_file_mounts(
 
     @timeline.event
     @usage_lib.messages.usage.update_runtime('setup')
-    def setup(self, handle: ResourceHandle, task: 'task_lib.Task') -> None:
-        return self._setup(handle, task)
+    def setup(self, handle: ResourceHandle, task: 'task_lib.Task',
+              async_setup: bool) -> None:
+        return self._setup(handle, task, async_setup)
 
     def add_storage_objects(self, task: 'task_lib.Task') -> None:
         raise NotImplementedError
 
     @timeline.event
     @usage_lib.messages.usage.update_runtime('execute')
-    def execute(self, handle: ResourceHandle, task: 'task_lib.Task',
-                detach_run: bool) -> None:
+    def execute(self,
+                handle: ResourceHandle,
+                task: 'task_lib.Task',
+                detach_run: bool,
+                setup_cmd: Optional[str] = None) -> None:
path: sky/backends/backend.py
text: @user1:
Just from reading the interface -- this seems odd as I'd expect we can read from `task.setup`?

@author:
I agree this is a bit odd. We now pass the `setup_cmd` through the member variable in the backend object instead.

diff_hunk: @@ -0,0 +1,27 @@
+# A task runs on an existing multinode cluster.
+#
+# Runs a task that requires multinode.
+#
+# Usage:
+#   sky launch -c mjq cluster_multinode.yaml
+#   sky exec mjq job_multinode.yaml
+#   sky exec mjq job.yaml
+
+name: job_multinode
+
path: examples/job_queue/job_multinode_long_setup.yaml
text: @user1:

nit: remove?

@author:
Done. Thanks!