Repository: skypilot-org/skypilot
Pull request: Add retry-until-up feature for launch and start
Closes #857.

Main changes:
1. Added retry until up for both `sky launch` and `sky start`.
2. Exponential backoff for retry in `--retry-until-up` and spot recovery.

Tested:
- [x] `sky launch -c test-gpu --gpus V100 ''`
- [x] `sky launch -c test-gpu --gpus V100 -r ''`
- [x] `sky launch -c retry-test --gpus=A100:8 --cloud aws '' -r` will keeps retrying.

Problem statement:
UX: add "wait until launched" option to `sky launch`
Currently, both @Michaelvll and Daniel have used/discussed using wrapping `sky launch` in a bash loop in case launching fails.  They used something like
```
 until sky launch ...; do; done
```
to manually make this work.  We should consider adding a `--wait-until-up / -w` or  `--retry-until-up / -r` option.

Hints from issue discussion:

Commit message: Add backoff for spot recovery

Patch to review:
diff --git a/sky/backends/backend.py b/sky/backends/backend.py
index d0274d19345..0d109b1e48e 100644
--- a/sky/backends/backend.py
+++ b/sky/backends/backend.py
@@ -28,7 +28,8 @@ def provision(self,
                   to_provision: Optional['resources.Resources'],
                   dryrun: bool,
                   stream_logs: bool,
-                  cluster_name: Optional[str] = None) -> ResourceHandle:
+                  cluster_name: Optional[str] = None,
+                  retry_until_up: bool = False) -> ResourceHandle:
         raise NotImplementedError
 
     def sync_workdir(self, handle: ResourceHandle, workdir: Path) -> None:
diff --git a/sky/backends/backend_utils.py b/sky/backends/backend_utils.py
index f71933ab800..98a953090f9 100644
--- a/sky/backends/backend_utils.py
+++ b/sky/backends/backend_utils.py
@@ -9,6 +9,7 @@
 import os
 import pathlib
 import psutil
+import random
 import re
 import shlex
 import socket
@@ -1441,3 +1442,30 @@ def stop_handler(signum, frame):
                    f'running after Ctrl-Z.{colorama.Style.RESET_ALL}')
     kill_children_processes()
     sys.exit(exceptions.SIGTSTP_CODE)
+
+
+class Backoff:
+    """Exponential backoff with jittering."""
+    MULTIPLIER = 1.6
+    JITTER = 0.4
+
+    def __init__(self, initial_backoff: int = 5, max_backoff_ratio: int = 20):
+        self._initial = True
+        self._backoff = None
+        self._inital_backoff = initial_backoff
+        self._max_backoff = max_backoff_ratio * self._inital_backoff
+
+    # https://github.com/grpc/grpc/blob/2d4f3c56001cd1e1f85734b2f7c5ce5f2797c38a/doc/connection-backoff.md
+    # https://github.com/grpc/grpc/blob/5fc3ff82032d0ebc4bf252a170ebe66aacf9ed9d/src/core/lib/backoff/backoff.cc
+
+    def current_backoff(self) -> float:
+        """Backs off once and returns the current backoff in seconds."""
+        if self._initial:
+            self._initial = False
+            self._backoff = min(self._inital_backoff, self._max_backoff)
+        else:
+            self._backoff = min(self._backoff * self.MULTIPLIER,
+                                self._max_backoff)
+        self._backoff += random.uniform(-self.JITTER * self._backoff,
+                                        self.JITTER * self._backoff)
+        return self._backoff
diff --git a/sky/backends/cloud_vm_ray_backend.py b/sky/backends/cloud_vm_ray_backend.py
index d777a55807a..37b9ea1f923 100644
--- a/sky/backends/cloud_vm_ray_backend.py
+++ b/sky/backends/cloud_vm_ray_backend.py
@@ -66,6 +66,13 @@
 _RSYNC_FILTER_OPTION = '--filter=\'dir-merge,- .gitignore\''
 _RSYNC_EXCLUDE_OPTION = '--exclude-from=.git/info/exclude'
 
+# Time gap between retries after provisioning on all possible place failed.
+_RETRY_UNTIL_UP_INIT_GAP_SECONDS = 60
+_RETRY_UNTIL_UP_MESSAGE = (f'{colorama.Style.BRIGHT}=== Retry until up ==='
+                           f'{colorama.Style.RESET_ALL}\n'
+                           'Retry until up is set, retry the provisioning '
+                           'after {gap_seconds:.1f}s...')
+
 
 def _get_cluster_config_template(cloud):
     cloud_to_template = {
@@ -1277,7 +1284,8 @@ def provision(self,
                   to_provision: Optional['resources_lib.Resources'],
                   dryrun: bool,
                   stream_logs: bool,
-                  cluster_name: Optional[str] = None) -> ResourceHandle:
+                  cluster_name: Optional[str] = None,
+                  retry_until_up: bool = False) -> ResourceHandle:
         """Provisions using 'ray up'."""
         # Try to launch the exiting cluster first
         if cluster_name is None:
@@ -1304,25 +1312,43 @@ def provision(self,
             # TODO(suquark): once we have sky on PYPI, we should directly
             # install sky from PYPI.
             local_wheel_path = wheel_utils.build_sky_wheel()
-            try:
-                provisioner = RetryingVmProvisioner(self.log_dir, self._dag,
-                                                    self._optimize_target,
-                                                    local_wheel_path)
-                config_dict = provisioner.provision_with_retries(
-                    task, to_provision_config, dryrun, stream_logs)
-            except exceptions.ResourcesUnavailableError as e:
-                # Do not remove the stopped cluster from the global state
-                # if failed to start.
-                if e.no_retry:
-                    logger.error(e)
+            backoff = backend_utils.Backoff(_RETRY_UNTIL_UP_INIT_GAP_SECONDS)
+            while True:
+                try:
+                    provisioner = RetryingVmProvisioner(self.log_dir, self._dag,
+                                                        self._optimize_target,
+                                                        local_wheel_path)
+                    config_dict = provisioner.provision_with_retries(
+                        task, to_provision_config, dryrun, stream_logs)
+                    break
+                except exceptions.ResourcesUnavailableError as e:
+                    # Do not remove the stopped cluster from the global state
+                    # if failed to start.
+                    if e.no_retry:
+                        logger.error(e)
+                        if retry_until_up:
+                            gap_seconds = backoff.current_backoff()
+                            logger.info(
+                                _RETRY_UNTIL_UP_MESSAGE.format(
+                                    gap_seconds=gap_seconds))
+                            time.sleep(gap_seconds)
+                            continue
+                        sys.exit(1)
+                    # Clean up the cluster's entry in `sky status`.
+                    global_user_state.remove_cluster(cluster_name,
+                                                     terminate=True)
+                    logger.error(
+                        'Failed to provision all possible launchable resources.'
+                        f' Relax the task\'s resource requirements:\n '
+                        f'{task.num_nodes}x {task.resources}')
+                    if retry_until_up:
+                        gap_seconds = backoff.current_backoff()
+                        logger.info(
+                            _RETRY_UNTIL_UP_MESSAGE.format(
+                                gap_seconds=gap_seconds))
+                        time.sleep(gap_seconds)
+                        continue
                     sys.exit(1)
-                # Clean up the cluster's entry in `sky status`.
-                global_user_state.remove_cluster(cluster_name, terminate=True)
-                logger.error(
-                    'Failed to provision all possible launchable resources. '
-                    f'Relax the task\'s resource requirements:\n '
-                    f'{task.num_nodes}x {task.resources}')
-                sys.exit(1)
             if dryrun:
                 return
             cluster_config_file = config_dict['ray']
diff --git a/sky/backends/local_docker_backend.py b/sky/backends/local_docker_backend.py
index 762e3296a80..ba6e129c603 100644
--- a/sky/backends/local_docker_backend.py
+++ b/sky/backends/local_docker_backend.py
@@ -146,15 +146,18 @@ def provision(self,
                   to_provision: Optional['resources.Resources'],
                   dryrun: bool,
                   stream_logs: bool,
-                  cluster_name: Optional[str] = None) -> ResourceHandle:
+                  cluster_name: Optional[str] = None,
+                  retry_until_up: bool = False) -> ResourceHandle:
         """
         Builds docker image for the task and returns the cluster name as handle.
 
         Since resource demands are ignored, There's no provisioning in local
         docker.
         """
-        assert task.name is not None, 'Task name cannot be None - have you ' \
-                                      'specified a task name?'
+        assert task.name is not None, ('Task name cannot be None - have you '
+                                       'specified a task name?')
+        assert not retry_until_up, ('Retrying until up is not supported in '
+                                    f'{self.NAME}')
         if cluster_name is None:
             cluster_name = backend_utils.generate_cluster_name()
         if stream_logs:
diff --git a/sky/cli.py b/sky/cli.py
index 3c13c26b71e..5b7e25cdb40 100644
--- a/sky/cli.py
+++ b/sky/cli.py
@@ -512,7 +512,8 @@ def _check_yaml(entrypoint: str) -> bool:
 
 
 def _start_cluster(cluster_name: str,
-                   idle_minutes_to_autostop: Optional[int] = None):
+                   idle_minutes_to_autostop: Optional[int] = None,
+                   retry_until_up: bool = False):
     handle = global_user_state.get_handle_from_cluster_name(cluster_name)
     backend = backend_utils.get_backend_from_handle(handle)
     assert isinstance(backend, backends.CloudVmRayBackend)
@@ -523,7 +524,8 @@ def _start_cluster(cluster_name: str,
                                to_provision=handle.launched_resources,
                                dryrun=False,
                                stream_logs=True,
-                               cluster_name=cluster_name)
+                               cluster_name=cluster_name,
+                               retry_until_up=retry_until_up)
     if idle_minutes_to_autostop is not None:
         backend.set_autostop(handle, idle_minutes_to_autostop)
     return handle
@@ -593,6 +595,14 @@ def cli():
           'of idleness after setup/file_mounts. This is equivalent to '
           'running `sky launch -d ...` and then `sky autostop -i <minutes>`. '
           'If not set, the cluster will not be auto-stopped.'))
+@click.option(
+    '--retry-until-up',
+    '-r',
+    default=False,
+    is_flag=True,
+    required=False,
+    help=('If True, retry until the cluster is up even after sky fails to '
+          'launch cluster on any possible region/cloud.'))
 @click.option('--yes',
               '-y',
               is_flag=True,
@@ -615,6 +625,7 @@ def launch(
     env: List[Dict[str, str]],
     disk_size: Optional[int],
     idle_minutes_to_autostop: Optional[int],
+    retry_until_up: bool,
     yes: bool,
 ):
     """Launch a task from a YAML or a command (rerun setup if cluster exists).
@@ -715,7 +726,8 @@ def launch(
                cluster_name=cluster,
                detach_run=detach_run,
                backend=backend,
-               idle_minutes_to_autostop=idle_minutes_to_autostop)
+               idle_minutes_to_autostop=idle_minutes_to_autostop,
+               retry_until_up=retry_until_up)
 
 
 @cli.command(cls=_DocumentedCodeCommand)
@@ -1223,7 +1235,13 @@ def autostop(
               default=False,
               required=False,
               help='Skip confirmation prompt.')
-def start(clusters: Tuple[str], yes: bool):
+@click.option('--retry-until-up',
+              '-r',
+              default=False,
+              is_flag=True,
+              required=False,
+              help='Retry until cluster is up.')
+def start(clusters: Tuple[str], yes: bool, retry_until_up: bool):
     """Restart cluster(s).
 
     If a cluster is previously stopped (status == STOPPED) or failed in
@@ -1312,7 +1330,7 @@ def start(clusters: Tuple[str], yes: bool):
             show_default=True)
 
     for name in to_start:
-        _start_cluster(name)
+        _start_cluster(name, retry_until_up=retry_until_up)
         click.secho(f'Cluster {name} started.', fg='green')
 
 
diff --git a/sky/execution.py b/sky/execution.py
index a7e4c3f5f7e..3c64929eee5 100644
--- a/sky/execution.py
+++ b/sky/execution.py
@@ -44,17 +44,20 @@ class Stage(enum.Enum):
     TEARDOWN = enum.auto()
 
 
-def _execute(dag: sky.Dag,
-             dryrun: bool = False,
-             teardown: bool = False,
-             stream_logs: bool = True,
-             handle: Any = None,
-             backend: Optional[backends.Backend] = None,
-             optimize_target: OptimizeTarget = OptimizeTarget.COST,
-             stages: Optional[List[Stage]] = None,
-             cluster_name: Optional[str] = None,
-             detach_run: bool = False,
-             idle_minutes_to_autostop: Optional[int] = None) -> None:
+def _execute(
+    dag: sky.Dag,
+    dryrun: bool = False,
+    teardown: bool = False,
+    stream_logs: bool = True,
+    handle: Any = None,
+    backend: Optional[backends.Backend] = None,
+    retry_until_up: bool = False,
+    optimize_target: OptimizeTarget = OptimizeTarget.COST,
+    stages: Optional[List[Stage]] = None,
+    cluster_name: Optional[str] = None,
+    detach_run: bool = False,
+    idle_minutes_to_autostop: Optional[int] = None,
+) -> None:
     """Runs a DAG.
 
     If the DAG has not been optimized yet, this will call sky.optimize() for
@@ -71,6 +74,8 @@ def _execute(dag: sky.Dag,
         handle instead of provisioning a new one.
       backend: Backend; backend to use for executing the tasks. Defaults to
         CloudVmRayBackend()
+      retry_until_up: bool; whether to retry the provisioning until the cluster
+        is up.
       optimize_target: OptimizeTarget; the dag optimization metric, e.g.
         OptimizeTarget.COST.
       stages: List of stages to run.  If None, run the whole life cycle of
@@ -131,7 +136,8 @@ def _execute(dag: sky.Dag,
                                            task.best_resources,
                                            dryrun=dryrun,
                                            stream_logs=stream_logs,
-                                           cluster_name=cluster_name)
+                                           cluster_name=cluster_name,
+                                           retry_until_up=retry_until_up)
 
         if dryrun:
             logger.info('Dry run finished.')
@@ -192,29 +198,35 @@ def _execute(dag: sky.Dag,
             print('\x1b[?25h', end='')  # Show cursor.
 
 
-def launch(dag: sky.Dag,
-           dryrun: bool = False,
-           teardown: bool = False,
-           stream_logs: bool = True,
-           backend: Optional[backends.Backend] = None,
-           optimize_target: OptimizeTarget = OptimizeTarget.COST,
-           cluster_name: Optional[str] = None,
-           detach_run: bool = False,
-           idle_minutes_to_autostop: Optional[int] = None,
-           is_spot_controller_task: bool = False) -> None:
+def launch(
+    dag: sky.Dag,
+    dryrun: bool = False,
+    teardown: bool = False,
+    stream_logs: bool = True,
+    backend: Optional[backends.Backend] = None,
+    retry_until_up: bool = False,
+    optimize_target: OptimizeTarget = OptimizeTarget.COST,
+    cluster_name: Optional[str] = None,
+    detach_run: bool = False,
+    idle_minutes_to_autostop: Optional[int] = None,
+    is_spot_controller_task: bool = False,
+) -> None:
     if not is_spot_controller_task:
         backend_utils.check_cluster_name_not_reserved(
             cluster_name, operation_str='sky.launch')
-    _execute(dag=dag,
-             dryrun=dryrun,
-             teardown=teardown,
-             stream_logs=stream_logs,
-             handle=None,
-             backend=backend,
-             optimize_target=optimize_target,
-             cluster_name=cluster_name,
-             detach_run=detach_run,
-             idle_minutes_to_autostop=idle_minutes_to_autostop)
+    _execute(
+        dag=dag,
+        dryrun=dryrun,
+        teardown=teardown,
+        stream_logs=stream_logs,
+        handle=None,
+        backend=backend,
+        retry_until_up=retry_until_up,
+        optimize_target=optimize_target,
+        cluster_name=cluster_name,
+        detach_run=detach_run,
+        idle_minutes_to_autostop=idle_minutes_to_autostop,
+    )
 
 
 def exec(  # pylint: disable=redefined-builtin
diff --git a/sky/spot/recovery_strategy.py b/sky/spot/recovery_strategy.py
index b1001123e64..1880fa1f06a 100644
--- a/sky/spot/recovery_strategy.py
+++ b/sky/spot/recovery_strategy.py
@@ -56,7 +56,7 @@ def make(cls, cluster_name: str, backend: 'backends.Backend',
 
     def launch(self,
                max_retry=3,
-               retry_gap_seconds=60,
+               retry_init_gap_seconds=60,
                raise_on_failure=True) -> Optional[float]:
         """Launch the spot cluster for the first time.
 
@@ -67,6 +67,7 @@ def launch(self,
         """
         # TODO(zhwu): handle the failure during `preparing sky runtime`.
         retry_cnt = 0
+        backoff = backend_utils.Backoff(retry_init_gap_seconds)
         while True:
             retry_cnt += 1
             try:
@@ -102,10 +103,11 @@ def launch(self,
                         'retries.')
                 else:
                     return None
+            gap_seconds = backoff.current_backoff()
             logger.info(
-                f'Retrying to launch the spot cluster in {retry_gap_seconds} '
+                f'Retrying to launch the spot cluster in {gap_seconds:.1f} '
                 'seconds.')
-            time.sleep(retry_gap_seconds)
+            time.sleep(gap_seconds)
 
     def recover(self) -> float:
         """Relaunch the spot cluster after failure and wait until job starts.
@@ -129,7 +131,7 @@ class FailoverStrategyExecutor(StrategyExecutor, name='FAILOVER', default=True):
     """Failover strategy: wait in same region and failover after timout."""
 
     _MAX_RETRY_CNT = 240  # Retry for 4 hours.
-    _RETRY_GAP_SECONDS = 60
+    _RETRY_INIT_GAP_SECONDS = 60
 
     def recover(self) -> float:
         # 1. Cancel the jobs and launch the cluster with the STOPPED status,
@@ -159,12 +161,13 @@ def recover(self) -> float:
         self.terminate_cluster()
 
         # Step 3
-        launched_time = self.launch(max_retry=self._MAX_RETRY_CNT,
-                                    retry_gap_seconds=self._RETRY_GAP_SECONDS,
-                                    raise_on_failure=False)
+        launched_time = self.launch(
+            max_retry=self._MAX_RETRY_CNT,
+            retry_init_gap_seconds=self._RETRY_INIT_GAP_SECONDS,
+            raise_on_failure=False)
         if launched_time is None:
             logger.error(f'Failed to recover the spot cluster after retrying '
                          f'{self._MAX_RETRY_CNT} times every '
-                         f'{self._RETRY_GAP_SECONDS} seconds.')
+                         f'{self._RETRY_INIT_GAP_SECONDS} seconds.')
             raise RuntimeError('Failed to recover the spot cluster.')
         return launched_time


Code review comments:
diff_hunk: @@ -159,12 +161,13 @@ def recover(self) -> float:
         self.terminate_cluster()
 
         # Step 3
-        launched_time = self.launch(max_retry=self._MAX_RETRY_CNT,
-                                    retry_gap_seconds=self._RETRY_GAP_SECONDS,
-                                    raise_on_failure=False)
+        launched_time = self.launch(
+            max_retry=self._MAX_RETRY_CNT,
+            retry_init_gap_seconds=self._RETRY_INIT_GAP_SECONDS,
+            raise_on_failure=False)
         if launched_time is None:
             logger.error(f'Failed to recover the spot cluster after retrying '
                          f'{self._MAX_RETRY_CNT} times every '
-                         f'{self._RETRY_GAP_SECONDS} seconds.')
+                         f'{self._RETRY_INIT_GAP_SECONDS} seconds.')
path: sky/spot/recovery_strategy.py
text: @user1:
Can delete "every ... seconds" now, as it's not accurate due to the exponential backoff.

@author:
Good point! Removed. Thanks!

diff_hunk: @@ -1304,25 +1312,43 @@ def provision(self,
             # TODO(suquark): once we have sky on PYPI, we should directly
             # install sky from PYPI.
             local_wheel_path = wheel_utils.build_sky_wheel()
-            try:
-                provisioner = RetryingVmProvisioner(self.log_dir, self._dag,
-                                                    self._optimize_target,
-                                                    local_wheel_path)
-                config_dict = provisioner.provision_with_retries(
-                    task, to_provision_config, dryrun, stream_logs)
-            except exceptions.ResourcesUnavailableError as e:
-                # Do not remove the stopped cluster from the global state
-                # if failed to start.
-                if e.no_retry:
-                    logger.error(e)
+            backoff = backend_utils.Backoff(_RETRY_UNTIL_UP_INIT_GAP_SECONDS)
path: sky/backends/cloud_vm_ray_backend.py
text: @user1:
How about
```
backoff = backend_utils.Backoff(initial_backoff=_RETRY_UNTIL_UP_INIT_GAP_SECONDS, max_backoff_ratio=5)
```
A max backoff of ~5 minutes seems better than ~20 mins.

@author:
Good point! Fixed by changing the default value of `max_backoff_ratio` to 5.

diff_hunk: @@ -1304,25 +1312,43 @@ def provision(self,
             # TODO(suquark): once we have sky on PYPI, we should directly
             # install sky from PYPI.
             local_wheel_path = wheel_utils.build_sky_wheel()
-            try:
-                provisioner = RetryingVmProvisioner(self.log_dir, self._dag,
-                                                    self._optimize_target,
-                                                    local_wheel_path)
-                config_dict = provisioner.provision_with_retries(
-                    task, to_provision_config, dryrun, stream_logs)
-            except exceptions.ResourcesUnavailableError as e:
-                # Do not remove the stopped cluster from the global state
-                # if failed to start.
-                if e.no_retry:
-                    logger.error(e)
+            backoff = backend_utils.Backoff(_RETRY_UNTIL_UP_INIT_GAP_SECONDS)
+            while True:
+                try:
+                    provisioner = RetryingVmProvisioner(self.log_dir, self._dag,
+                                                        self._optimize_target,
+                                                        local_wheel_path)
+                    config_dict = provisioner.provision_with_retries(
+                        task, to_provision_config, dryrun, stream_logs)
+                    break
+                except exceptions.ResourcesUnavailableError as e:
+                    # Do not remove the stopped cluster from the global state
+                    # if failed to start.
+                    if e.no_retry:
+                        logger.error(e)
+                        if retry_until_up:
path: sky/backends/cloud_vm_ray_backend.py
text: @user1:
It's a bit confusing that we are under `e.no_retry` while this block is actually retrying.  How about we rename that field to `no_failover`?

@user1:
Also, consider factoring out these lines -- also used for L1345-1349 -- into a short helper func under L1324.

@author:
Good point! Instead of having a function, I moved the `retry_until_up` to the same place for both condition to reduce the code duplication.

diff_hunk: @@ -146,15 +146,18 @@ def provision(self,
                   to_provision: Optional['resources.Resources'],
                   dryrun: bool,
                   stream_logs: bool,
-                  cluster_name: Optional[str] = None) -> ResourceHandle:
+                  cluster_name: Optional[str] = None,
+                  retry_until_up: bool = False) -> ResourceHandle:
         """
         Builds docker image for the task and returns the cluster name as handle.
 
         Since resource demands are ignored, There's no provisioning in local
         docker.
         """
-        assert task.name is not None, 'Task name cannot be None - have you ' \
-                                      'specified a task name?'
+        assert task.name is not None, ('Task name cannot be None - have you '
+                                       'specified a task name?')
+        assert not retry_until_up, ('Retrying until up is not supported in '
+                                    f'{self.NAME}')
path: sky/backends/local_docker_backend.py
text: @user1:
Probably ok to print a message saying ignoring this flag.  Useful for e.g., user uses the same command for cloud and then adds a --docker for exploration.

@author:
Good point. Fixed. Thanks!

diff_hunk: @@ -66,6 +66,13 @@
 _RSYNC_FILTER_OPTION = '--filter=\'dir-merge,- .gitignore\''
 _RSYNC_EXCLUDE_OPTION = '--exclude-from=.git/info/exclude'
 
+# Time gap between retries after provisioning on all possible place failed.
+_RETRY_UNTIL_UP_INIT_GAP_SECONDS = 60
+_RETRY_UNTIL_UP_MESSAGE = (f'{colorama.Style.BRIGHT}=== Retry until up ==='
+                           f'{colorama.Style.RESET_ALL}\n'
+                           'Retry until up is set, retry the provisioning '
path: sky/backends/cloud_vm_ray_backend.py
text: @user1:
nit: can be shorter as L71 already hints that the flag is set, `f'Retrying provisioning after {gap_seconds:.0f}s.'`

@user1:
After trying this out on V100:8 spot for a long time, I think our UX can be improved by
- showing an integer seconds
- we can show `f'Retrying provisioning after {gap_seconds:.0f}s (exponential backoff with random jittering). Already tried N attempts.'` where N starts with 0

@author:
Good point! Fixed. Thanks!

@author:
The attempt may start from 1, since the message will be shown after 1 failure among all regions and clouds.