Repository: skypilot-org/skypilot
Pull request: [Core] Fix status refresh for INIT cluster
<!-- Describe the changes in this PR -->

Fixes #2490 

<!-- 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):

- [x] Code formatting: `bash format.sh`
- [x] Any manual or new tests for this PR (please specify below)
  - [x] reproducible code in #2490
- [ ] 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:
[Core] Assertion error when refreshing status for INIT cluster
To reproduce:
```
sky launch -c existing-up-cluster
# ctrl-c when it is launching
sky status -r existing-up-cluster
sky launch -c existing-up-cluster
```

```
sky/backends/backend_utils.py", line 2024, in run_ray_status_to_check_ray_cluster_healthy
    assert handle.head_ssh_port is not None, handle
AssertionError: ResourceHandle(
        cluster_name=spot-dev-zhwu,
        cluster_name_on_cloud=spot-dev-zhwu,
        head_ip=34.135.151.192,
        stable_internal_external_ips=[('10.128.0.68', '34.135.151.192')],
        stable_ssh_ports=None,
        cluster_yaml=/Users/zhwu/.sky/generated/spot-dev-zhwu.yml, 
        launched_resources=1x GCP(n2-standard-8, image_id={'us-central1': 'projects/skypilot-managed-spot/global/machineImages/spot-dev-image'}), 
        docker_user=None,
        tpu_create_script=None, 
        tpu_delete_script=None)
```

**This is a serious problem, as once the cluster enter this state, none of our clis can operate on the cluster to get it out of the state.**

Hints from issue discussion:

Commit message: changed to runtime error

Patch to review:
diff --git a/sky/backends/backend_utils.py b/sky/backends/backend_utils.py
index 64564f3eb0d..410e3df9a1c 100644
--- a/sky/backends/backend_utils.py
+++ b/sky/backends/backend_utils.py
@@ -2014,14 +2014,22 @@ def run_ray_status_to_check_ray_cluster_healthy() -> bool:
             # triggered.
             if external_ips is None or len(external_ips) == 0:
                 logger.debug(f'Refreshing status ({cluster_name!r}): No cached '
-                             f'IPs found. External IPs: {external_ips}')
+                             f'IPs found. Handle: {handle}')
                 raise exceptions.FetchIPError(
                     reason=exceptions.FetchIPError.Reason.HEAD)
 
+            if handle.head_ssh_port is None:
+                # Refresh the ssh ports. It is ok to refresh as it is fast.
+                handle.external_ssh_ports()
+                if handle.head_ssh_port is None:
+                    raise RuntimeError(
+                        f'Refreshing status ({cluster_name!r}): failed '
+                        f'to get the ssh ports. Handle: {handle}')
+
             # Check if ray cluster status is healthy.
             ssh_credentials = ssh_credential_from_yaml(handle.cluster_yaml,
                                                        handle.docker_user)
-            assert handle.head_ssh_port is not None, handle
+
             runner = command_runner.SSHCommandRunner(external_ips[0],
                                                      **ssh_credentials,
                                                      port=handle.head_ssh_port)
diff --git a/sky/backends/cloud_vm_ray_backend.py b/sky/backends/cloud_vm_ray_backend.py
index 1e99b6f4ad1..2f7e8fe9d23 100644
--- a/sky/backends/cloud_vm_ray_backend.py
+++ b/sky/backends/cloud_vm_ray_backend.py
@@ -1524,9 +1524,10 @@ def _retry_zones(
             if zones and len(zones) == 1:
                 launched_resources = launched_resources.copy(zone=zones[0].name)
 
-            prev_cluster_ips = None
+            prev_cluster_ips, prev_ssh_ports = None, None
             if prev_handle is not None:
                 prev_cluster_ips = prev_handle.stable_internal_external_ips
+                prev_ssh_ports = prev_handle.stable_ssh_ports
             # Record early, so if anything goes wrong, 'sky status' will show
             # the cluster name and users can appropriately 'sky down'.  It also
             # means a second 'sky launch -c <name>' will attempt to reuse.
@@ -1543,10 +1544,11 @@ def _retry_zones(
                 launched_resources=launched_resources,
                 tpu_create_script=config_dict.get('tpu-create-script'),
                 tpu_delete_script=config_dict.get('tpu-delete-script'),
-                # Use the previous cluster's IPs if available to optimize
-                # the case where the cluster is restarted, i.e., no need to
-                # query the IPs from the cloud provider.
-                stable_internal_external_ips=prev_cluster_ips)
+                # Use the previous cluster's IPs and ports if available to
+                # optimize the case where the cluster is restarted, i.e., no
+                # need to query the IPs from the cloud provider.
+                stable_internal_external_ips=prev_cluster_ips,
+                stable_ssh_ports=prev_ssh_ports)
             usage_lib.messages.usage.update_final_cluster_status(
                 status_lib.ClusterStatus.INIT)
 


Code review comments:
diff_hunk: @@ -2014,14 +2014,22 @@ def run_ray_status_to_check_ray_cluster_healthy() -> bool:
             # triggered.
             if external_ips is None or len(external_ips) == 0:
                 logger.debug(f'Refreshing status ({cluster_name!r}): No cached '
-                             f'IPs found. External IPs: {external_ips}')
+                             f'IPs found. Handle: {handle}')
                 raise exceptions.FetchIPError(
                     reason=exceptions.FetchIPError.Reason.HEAD)
 
+            if handle.head_ssh_port is None:
+                # Refresh the ssh ports. It is ok to refresh as it is fast.
+                handle.external_ssh_ports()
path: sky/backends/backend_utils.py
text: @user1:
Shall we use `update_ssh_ports` for better readability?

@author:
Good point! Actually, we don't need to handle this `None` or refresh case, as we already have it in the `external_ssh_ports` call. changed to that instead.