Retry for azure when first launched (skypilot-org#985)

* retry for azure head timeout

* check more matching

* Add retries

* reformat

* increase retries for node ip query

* format

* Address comments

* Fix readme for env

* Address comments

* format

Author: Michaelvll

README.md

@@ -32,7 +32,8 @@ IMPORTANT: Please `export SKY_DEV=1` before running the sky commands in the term
 
 
 ### Environment Variable Options
-- `export SKY_DEV=1` to show debugging logs (logging.DEBUG) and send the logs to dev space.
+- `export SKY_DEV=1` to send the logs to dev space.
+- `export SKY_DEBUG=1` to show debugging logs (logging.DEBUG).
 - `export SKY_DISABLE_USAGE_COLLECTION=1` to disable usage logging.
 - `export SKY_MINIMIZE_LOGGING=1` to minimize the sky outputs for demo purpose.
 

sky/backends/backend_utils.py

@@ -925,6 +925,7 @@ def generate_cluster_name():
 
 def query_head_ip_with_retries(cluster_yaml: str, max_attempts: int = 1) -> str:
     """Returns the ip of the head node from yaml file."""
+    backoff = Backoff(initial_backoff=5, max_backoff_factor=5)
     for i in range(max_attempts):
         try:
             out = subprocess_utils.run(
@@ -940,7 +941,7 @@ def query_head_ip_with_retries(cluster_yaml: str, max_attempts: int = 1) -> str:
                 raise RuntimeError('Failed to get head ip') from e
             # Retry if the cluster is not up yet.
             logger.debug('Retrying to get head ip.')
-            time.sleep(5)
+            time.sleep(backoff.current_backoff())
     return head_ip
 
 

sky/backends/cloud_vm_ray_backend.py

@@ -68,7 +68,7 @@
 _RETRY_UNTIL_UP_INIT_GAP_SECONDS = 60
 
 # The maximum retry count for fetching head IP address.
-_HEAD_IP_MAX_ATTEMPTS = 3
+_HEAD_IP_MAX_ATTEMPTS = 5
 
 _TEARDOWN_FAILURE_MESSAGE = (
     f'{colorama.Fore.RED}Failed to terminate '
@@ -89,6 +89,8 @@
 
 _TPU_NOT_FOUND_ERROR = 'ERROR: (gcloud.compute.tpus.delete) NOT_FOUND'
 
+_MAX_RAY_UP_RETRY = 5
+
 
 def _get_cluster_config_template(cloud):
     cloud_to_template = {
@@ -1058,18 +1060,51 @@ def ray_up():
         # Edge case: /tmp/ray does not exist, so autoscaler can't create/store
         # cluster lock and cluster state.
         os.makedirs('/tmp/ray', exist_ok=True)
-        returncode, stdout, stderr = ray_up()
-        if (returncode != 0 and 'Processing file mounts' in stdout and
-                'Running setup commands' not in stdout):
-            # Retry ray up if it failed due to file mounts, because it is
-            # probably due to too many ssh connections issue and can be fixed
-            # by retrying.
-            # This is required when using custom image for GCP.
-            logger.info(
-                'Retrying sky runtime setup due to ssh connection issue.')
-            returncode, stdout, stderr = ray_up()
 
-        logger.debug(f'Ray up takes {time.time() - start} seconds.')
+        # Launch the cluster with ray up
+
+        # Retry if the any of the following happens:
+        # 1. Failed due to timeout when fetching head node for Azure.
+        # 2. Failed due to file mounts, because it is probably has too
+        # many ssh connections and can be fixed by retrying.
+        # This is required when using custom image for GCP.
+        def need_ray_up(
+                ray_up_return_value: Optional[Tuple[int, str, str]]) -> bool:
+
+            # Indicates the first ray up.
+            if ray_up_return_value is None:
+                return True
+
+            returncode, stdout, stderr = ray_up_return_value
+            if returncode == 0:
+                return False
+
+            if ('Head node fetch timed out. Failed to create head node.'
+                    in stderr and isinstance(to_provision_cloud, clouds.Azure)):
+                logger.info(
+                    'Retrying head node provisioning due to head fetching '
+                    'timeout.')
+                return True
+            if ('Processing file mounts' in stdout and
+                    'Running setup commands' not in stdout and
+                    'Failed to setup head node.' in stderr):
+                logger.info(
+                    'Retrying sky runtime setup due to ssh connection issue.')
+                return True
+            return False
+
+        retry_cnt = 0
+        ray_up_return_value = None
+        while (retry_cnt < _MAX_RAY_UP_RETRY and
+               need_ray_up(ray_up_return_value)):
+            retry_cnt += 1
+            ray_up_return_value = ray_up()
+
+        assert ray_up_return_value is not None
+        returncode, stdout, stderr = ray_up_return_value
+
+        logger.debug(f'Ray up takes {time.time() - start} seconds with '
+                     f'{retry_cnt} retries.')
 
         # Only 1 node or head node provisioning failure.
         if num_nodes == 1 and returncode == 0:
@@ -1787,9 +1822,14 @@ def _exec_code_on_head(
                 f'--address=127.0.0.1:8265 --job-id {ray_job_id} --no-wait '
                 f'-- "{executable} -u {script_path} > {remote_log_path} 2>&1"')
 
-        returncode = self.run_on_head(handle, job_submit_cmd, stream_logs=False)
-        subprocess_utils.handle_returncode(returncode, job_submit_cmd,
-                                           f'Failed to submit job {job_id}.')
+        returncode, stdout, stderr = self.run_on_head(handle,
+                                                      job_submit_cmd,
+                                                      stream_logs=False,
+                                                      require_outputs=True)
+        subprocess_utils.handle_returncode(returncode,
+                                           job_submit_cmd,
+                                           f'Failed to submit job {job_id}.',
+                                           stderr=stdout + stderr)
 
         logger.info('Job submitted with Job ID: '
                     f'{style.BRIGHT}{job_id}{style.RESET_ALL}')