Repository: skypilot-org/skypilot
Pull request: [Core] Wait `ray status` to be ready
<!-- Describe the changes in this PR -->

Fixes #3129

This PR also raises proper error for the case when Azure have an invalid state string. Fixes #3134

<!-- Describe the tests ran -->
<!-- Unit tests (tests/test_*.py) are part of GitHub CI; below are tests that launch on the cloud. -->

Tested (run the relevant ones):

- [ ] Code formatting: `bash format.sh`
- [ ] Any manual or new tests for this PR (please specify below)
  - [x] reproducible code in #3129
  - [x] `time sky launch -c test-aws --cloud aws -y` (time for waiting ray cluster to start: 6s, this PR: 11s (wait for 5 seconds))
- [ ] All smoke tests: `pytest tests/test_smoke.py` 
- [ ] Relevant individual smoke tests: `pytest tests/test_smoke.py::test_fill_in_the_name` 
- [ ] Backward compatibility tests: `bash tests/backward_comaptibility_tests.sh`


Problem statement:
[Spot] Spot jobs on Azure always says it is preempted before submitting the job
A user reported that a managed spot job on Azure always says it is preempted after the cluster is launched before submitting the job.

This is likely because the `ray cluster` takes longer than other clouds to initalize before `ray status` can return the correct value of the cluster.

[Azure] Invalid state string when querying the instance status
A user reported an issue when running `sky serve` with spot replicas, where they got an empty string from the azure cli.
```
E 02-08 21:57:24 replica_managers.py:1000]   File "/home/azureuser/miniconda3/lib/python3.10/site-packages/sky/backends/backend_utils.py", line 1765, in _update_cluster_status_no_lock
E 02-08 21:57:24 replica_managers.py:1000]     node_statuses = _query_cluster_status_via_cloud_api(handle)
E 02-08 21:57:24 replica_managers.py:1000]   File "/home/azureuser/miniconda3/lib/python3.10/site-packages/sky/backends/backend_utils.py", line 1630, in _query_cluster_status_via_cloud_api
E 02-08 21:57:24 replica_managers.py:1000]     node_statuses = cloud.query_status(
E 02-08 21:57:24 replica_managers.py:1000]   File "/home/azureuser/miniconda3/lib/python3.10/site-packages/sky/clouds/azure.py", line 655, in query_status
E 02-08 21:57:24 replica_managers.py:1000]     node_status = status_map[s]
E 02-08 21:57:24 replica_managers.py:1000] KeyError: ''
```

Hints from issue discussion:
To reproduce:
`sky launch --cpus 2+ --cloud azure -c test-azure-query; sky status -r test-azure-query`

Commit message: add comment

Patch to review:
diff --git a/sky/backends/backend_utils.py b/sky/backends/backend_utils.py
index 248ddb1dc17..55666ea1d3f 100644
--- a/sky/backends/backend_utils.py
+++ b/sky/backends/backend_utils.py
@@ -882,6 +882,10 @@ def write_cluster_config(
                 'dump_port_command': dump_port_command,
                 # Ray version.
                 'ray_version': constants.SKY_REMOTE_RAY_VERSION,
+                # Command for waiting ray cluster to be ready on head.
+                'ray_head_wait_initalized_command':
+                    instance_setup.RAY_HEAD_WAIT_INITALIZED_COMMAND,
+
                 # Cloud credentials for cloud storage.
                 'credentials': credentials,
                 # Sky remote utils.
diff --git a/sky/clouds/azure.py b/sky/clouds/azure.py
index 7a46a0d69a6..c757127606f 100644
--- a/sky/clouds/azure.py
+++ b/sky/clouds/azure.py
@@ -648,6 +648,9 @@ def query_status(cls, name: str, tag_filters: Dict[str, str],
             original_statuses_list = [original_statuses_list]
         statuses = []
         for s in original_statuses_list:
+            if s not in status_map:
+                raise exceptions.ClusterStatusFetchingError(
+                    f'Failed to parse status from Azure response: {stdout}')
             node_status = status_map[s]
             if node_status is not None:
                 statuses.append(node_status)
diff --git a/sky/provision/instance_setup.py b/sky/provision/instance_setup.py
index 8e215b79288..1cef98bae40 100644
--- a/sky/provision/instance_setup.py
+++ b/sky/provision/instance_setup.py
@@ -34,7 +34,8 @@
 _DUMP_RAY_PORTS = (
     'python -c \'import json, os; '
     f'json.dump({constants.SKY_REMOTE_RAY_PORT_DICT_STR}, '
-    f'open(os.path.expanduser("{constants.SKY_REMOTE_RAY_PORT_FILE}"), "w"))\'')
+    f'open(os.path.expanduser("{constants.SKY_REMOTE_RAY_PORT_FILE}"), "w"))\';'
+)
 
 _RAY_PORT_COMMAND = (
     'RAY_PORT=$(python -c "from sky.skylet import job_lib; '
@@ -47,6 +48,15 @@
     f'{_RAY_PORT_COMMAND}; '
     'RAY_ADDRESS=127.0.0.1:$RAY_PORT ray status')
 
+# Command that waits for the ray status to be initialized. Otherwise, a later
+# `sky status -r` may fail due to the ray cluster not being ready.
+RAY_HEAD_WAIT_INITALIZED_COMMAND = (
+    f'while `RAY_ADDRESS=127.0.0.1:{constants.SKY_REMOTE_RAY_PORT} '
+    'ray status | grep -q "No cluster status."`; do '
+    'sleep 0.5; '
+    'echo "Waiting ray cluster to be initialized"; '
+    'done;')
+
 # Restart skylet when the version does not match to keep the skylet up-to-date.
 _MAYBE_SKYLET_RESTART_CMD = 'python3 -m sky.skylet.attempt_skylet'
 
@@ -243,7 +253,7 @@ def start_ray_on_head_node(cluster_name: str, custom_resource: Optional[str],
     cmd = ('ray stop; unset AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY; '
            'RAY_SCHEDULER_EVENTS=0 RAY_DEDUP_LOGS=0 '
            f'ray start --head {ray_options} || exit 1;' + _RAY_PRLIMIT +
-           _DUMP_RAY_PORTS)
+           _DUMP_RAY_PORTS + RAY_HEAD_WAIT_INITALIZED_COMMAND)
     logger.info(f'Running command on head node: {cmd}')
     # TODO(zhwu): add the output to log files.
     returncode, stdout, stderr = ssh_runner.run(cmd,
diff --git a/sky/templates/azure-ray.yml.j2 b/sky/templates/azure-ray.yml.j2
index 083d5463e79..856b3b86b67 100644
--- a/sky/templates/azure-ray.yml.j2
+++ b/sky/templates/azure-ray.yml.j2
@@ -169,9 +169,7 @@ setup_commands:
 #   current num items (num SSH connections): 2
 head_start_ray_commands:
   # NOTE: --disable-usage-stats in `ray start` saves 10 seconds of idle wait.
-  - ray stop; RAY_SCHEDULER_EVENTS=0 RAY_DEDUP_LOGS=0 ray start --disable-usage-stats --head --port={{ray_port}} --dashboard-port={{ray_dashboard_port}} --object-manager-port=8076 --autoscaling-config=~/ray_bootstrap_config.yaml {{"--resources='%s'" % custom_resources if custom_resources}} --temp-dir {{ray_temp_dir}} || exit 1;
-    which prlimit && for id in $(pgrep -f raylet/raylet); do sudo prlimit --nofile=1048576:1048576 --pid=$id || true; done;
-    {{dump_port_command}};
+  - ray stop; RAY_SCHEDULER_EVENTS=0 RAY_DEDUP_LOGS=0 ray start --disable-usage-stats --head --port={{ray_port}} --dashboard-port={{ray_dashboard_port}} --object-manager-port=8076 --autoscaling-config=~/ray_bootstrap_config.yaml {{"--resources='%s'" % custom_resources if custom_resources}} --temp-dir {{ray_temp_dir}} || exit 1; which prlimit && for id in $(pgrep -f raylet/raylet); do sudo prlimit --nofile=1048576:1048576 --pid=$id || true; done; {{dump_port_command}}; {{ray_head_wait_initalized_command}}
 
 {%- if num_nodes > 1 %}
 worker_start_ray_commands:
diff --git a/sky/templates/ibm-ray.yml.j2 b/sky/templates/ibm-ray.yml.j2
index 22ace28daba..4b8ad41962b 100644
--- a/sky/templates/ibm-ray.yml.j2
+++ b/sky/templates/ibm-ray.yml.j2
@@ -126,7 +126,7 @@ head_start_ray_commands:
   # all the sessions to be reloaded. This is a workaround.
   - ray stop; RAY_SCHEDULER_EVENTS=0 RAY_DEDUP_LOGS=0 ray start --disable-usage-stats --head --port={{ray_port}} --dashboard-port={{ray_dashboard_port}} --object-manager-port=8076 --autoscaling-config=~/ray_bootstrap_config.yaml {{"--resources='%s'" % custom_resources if custom_resources}} --temp-dir {{ray_temp_dir}} || exit 1;
     which prlimit && for id in $(pgrep -f raylet/raylet); do sudo prlimit --nofile=1048576:1048576 --pid=$id || true; done;
-    {{dump_port_command}};
+    {{dump_port_command}}; {{ray_head_wait_initalized_command}}
 
 {%- if num_nodes > 1 %}
 worker_start_ray_commands:
diff --git a/sky/templates/lambda-ray.yml.j2 b/sky/templates/lambda-ray.yml.j2
index 1250cb34334..942c9e388dc 100644
--- a/sky/templates/lambda-ray.yml.j2
+++ b/sky/templates/lambda-ray.yml.j2
@@ -96,7 +96,7 @@ setup_commands:
 head_start_ray_commands:
   - ray stop; RAY_SCHEDULER_EVENTS=0 RAY_DEDUP_LOGS=0 ray start --disable-usage-stats --head --port={{ray_port}} --dashboard-port={{ray_dashboard_port}} --object-manager-port=8076 --autoscaling-config=~/ray_bootstrap_config.yaml {{"--resources='%s'" % custom_resources if custom_resources}} --temp-dir {{ray_temp_dir}} || exit 1;
     which prlimit && for id in $(pgrep -f raylet/raylet); do sudo prlimit --nofile=1048576:1048576 --pid=$id || true; done;
-    {{dump_port_command}};
+    {{dump_port_command}}; {{ray_head_wait_initalized_command}}
 
 {%- if num_nodes > 1 %}
 worker_start_ray_commands:
diff --git a/sky/templates/oci-ray.yml.j2 b/sky/templates/oci-ray.yml.j2
index 15980bfa278..cf552d47e06 100644
--- a/sky/templates/oci-ray.yml.j2
+++ b/sky/templates/oci-ray.yml.j2
@@ -121,7 +121,7 @@ head_start_ray_commands:
   # all the sessions to be reloaded. This is a workaround.
   - ray stop; RAY_SCHEDULER_EVENTS=0 RAY_DEDUP_LOGS=0 ray start --disable-usage-stats --head --port={{ray_port}} --dashboard-port={{ray_dashboard_port}} --object-manager-port=8076 --autoscaling-config=~/ray_bootstrap_config.yaml {{"--resources='%s'" % custom_resources if custom_resources}} --temp-dir {{ray_temp_dir}} || exit 1;
     which prlimit && for id in $(pgrep -f raylet/raylet); do sudo prlimit --nofile=1048576:1048576 --pid=$id || true; done;
-    {{dump_port_command}};
+    {{dump_port_command}}; {{ray_head_wait_initalized_command}}
 
 {%- if num_nodes > 1 %}
 worker_start_ray_commands:
diff --git a/sky/templates/scp-ray.yml.j2 b/sky/templates/scp-ray.yml.j2
index ace20b70c8d..a0f1190761f 100644
--- a/sky/templates/scp-ray.yml.j2
+++ b/sky/templates/scp-ray.yml.j2
@@ -96,7 +96,7 @@ head_start_ray_commands:
   # all the sessions to be reloaded. This is a workaround.
   - ray stop; RAY_SCHEDULER_EVENTS=0 RAY_DEDUP_LOGS=0 ray start --disable-usage-stats --head --port={{ray_port}} --dashboard-port={{ray_dashboard_port}} --object-manager-port=8076 --autoscaling-config=~/ray_bootstrap_config.yaml {{"--resources='%s'" % custom_resources if custom_resources}} --temp-dir {{ray_temp_dir}} || exit 1;
     which prlimit && for id in $(pgrep -f raylet/raylet); do sudo prlimit --nofile=1048576:1048576 --pid=$id || true; done;
-    {{dump_port_command}};
+    {{dump_port_command}}; {{ray_head_wait_initalized_command}}
 
 {%- if num_nodes > 1 %}
 worker_start_ray_commands:


Code review comments:
diff_hunk: @@ -169,9 +169,7 @@ setup_commands:
 #   current num items (num SSH connections): 2
 head_start_ray_commands:
   # NOTE: --disable-usage-stats in `ray start` saves 10 seconds of idle wait.
-  - ray stop; RAY_SCHEDULER_EVENTS=0 RAY_DEDUP_LOGS=0 ray start --disable-usage-stats --head --port={{ray_port}} --dashboard-port={{ray_dashboard_port}} --object-manager-port=8076 --autoscaling-config=~/ray_bootstrap_config.yaml {{"--resources='%s'" % custom_resources if custom_resources}} --temp-dir {{ray_temp_dir}} || exit 1;
path: sky/templates/azure-ray.yml.j2
text: @user1:
Just curious: is there any reason to squash the original three lines into one?

@author:
Good catch! Changed it back.

diff_hunk: @@ -882,6 +882,10 @@ def write_cluster_config(
                 'dump_port_command': dump_port_command,
                 # Ray version.
                 'ray_version': constants.SKY_REMOTE_RAY_VERSION,
+                # Command for waiting ray cluster to be ready on head.
+                'ray_head_wait_initalized_command':
+                    instance_setup.RAY_HEAD_WAIT_INITALIZED_COMMAND,
path: sky/backends/backend_utils.py
text: @user1:
```suggestion
                'ray_head_wait_initialized_command':
                    instance_setup.RAY_HEAD_WAIT_INITIALIZED_COMMAND,
```
typo?

@author:
Good catch! Replaced all the occurance. Thanks!