Repository: skypilot-org/skypilot
Pull request: [Spot] Fix race condition for spot logs
This PR attempts to fix #1328.

The problem could happen due to a race condition, where the cluster is removed due to the preemption, but the `spot_status` has not been set to the correct state yet. In that case, the `handle` will be `None` and fails the log tailing.

i.e., after the following lines the cluster can be already removed from the cluster table, but the `spot_status` has not been set to `RECOVERING` yet.
https://github.com/skypilot-org/skypilot/blob/a66aa263325cec34f28c36a7f16f7ca1b23d130d/sky/spot/controller.py#L101-L107

Tested:
- [x] `sky spot launch -n test-logs "echo hi; sleep 200"` and terminate the cluster manually. (it may not be exactly the same as the case in #1328, where the preemption may happen before the logging starts)

Problem statement:
[Spot] `NoneType` for handle when tailing the logs for the spot job
@pounde reported an issue with the main branch, where the following issue occurs:
```
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/home/ubuntu/.local/lib/python3.9/site-packages/sky/spot/spot_utils.py", line 242, in stream_logs_by_id
    returncode = backend.tail_logs(handle,
  File "/home/ubuntu/.local/lib/python3.9/site-packages/sky/backends/cloud_vm_ray_backend.py", line 2390, in tail_logs
    job_owner = onprem_utils.get_job_owner(handle.cluster_yaml)
AttributeError: 'NoneType' object has no attribute 'cluster_yaml'
```

Hints from issue discussion:

Commit message: address comments

Patch to review:
diff --git a/sky/spot/spot_utils.py b/sky/spot/spot_utils.py
index 5dc402e1ba1..a3ffbfc7b0f 100644
--- a/sky/spot/spot_utils.py
+++ b/sky/spot/spot_utils.py
@@ -230,15 +230,21 @@ def stream_logs_by_id(job_id: int, follow: bool = True) -> str:
     cluster_name = generate_spot_cluster_name(task_name, job_id)
     backend = backends.CloudVmRayBackend()
     spot_status = spot_state.get_status(job_id)
-    while not spot_status.is_terminal():
-        if spot_status != spot_state.SpotStatus.RUNNING:
-            logger.info(f'INFO: The log is not ready yet, as the spot job '
-                        f'is {spot_status.value}. '
+    while spot_status is None or not spot_status.is_terminal():
+        handle = global_user_state.get_handle_from_cluster_name(cluster_name)
+        # Check the handle: The cluster can be removed from the table before the
+        # spot state is updated by the controller. In this case, we should skip
+        # the logging, and wait for the next round of status check.
+        if handle is None or spot_status != spot_state.SpotStatus.RUNNING:
+            status_help_str = ''
+            if (spot_status is not None and
+                    spot_status != spot_state.SpotStatus.RUNNING):
+                status_help_str = f', as the spot job is {spot_status.value}'
+            logger.info(f'INFO: The log is not ready yet{status_help_str}. '
                         f'Waiting for {JOB_STATUS_CHECK_GAP_SECONDS} seconds.')
             time.sleep(JOB_STATUS_CHECK_GAP_SECONDS)
             spot_status = spot_state.get_status(job_id)
             continue
-        handle = global_user_state.get_handle_from_cluster_name(cluster_name)
         returncode = backend.tail_logs(handle,
                                        job_id=None,
                                        spot_job_id=job_id,


Code review comments:
diff_hunk: @@ -230,15 +230,21 @@ def stream_logs_by_id(job_id: int, follow: bool = True) -> str:
     cluster_name = generate_spot_cluster_name(task_name, job_id)
     backend = backends.CloudVmRayBackend()
     spot_status = spot_state.get_status(job_id)
-    while not spot_status.is_terminal():
-        if spot_status != spot_state.SpotStatus.RUNNING:
-            logger.info(f'INFO: The log is not ready yet, as the spot job '
-                        f'is {spot_status.value}. '
+    while spot_status is None or not spot_status.is_terminal():
+        handle = global_user_state.get_handle_from_cluster_name(cluster_name)
+        # Check the handle: The cluster can be removed from the table before the
path: sky/spot/spot_utils.py
text: @user1:
```suggestion
        # Check the handle: The cluster can be preempted and removed from the table before the
```

@author:
Done. Thanks!