Repository: skypilot-org/skypilot
Pull request: Fix spot recovery caused by the exception raising
Closes #955.

After discussing with @infwinston, we found out that the problem is caused by #928, as all of our exception catching in the spot controller use `SystemExit`, which does not correctly catch the exception raised.

Tested:
- [x] `sky spot launch -n test-spot 'pip install tqdm; python -u -c "from tqdm import tqdm; import time; [(time.sleep(1), print(i)) for i in tqdm(range(1000000))]"'` and ctrl-c the log

Problem statement:
Potential bug of spot recovery 
I tried to run our spot training example by
```
sky spot launch examples/spot/bert_qa.yaml
```
and manually killed the VM after the training started to test our spot recovery.

But it was not successful and I got the `FAILED_CONTROLLER` error:
```
(sky-b7d3-weichiang pid=120832) E 07-06 19:18:16 controller.py:137] Unexpected error occurred: CommandError: Command python3 -u -c 'from sky.skylet import job_lib, log_lib;job_id = None if None is not None else job_lib.get_latest_job_id();job_status = job_lib.get_status(job_id);status_str = None if job_status is None else job_status.value;print("Job", job_id, status_str, flush=True)' failed with return code 255.
(sky-b7d3-weichiang pid=120832) E 07-06 19:18:16 controller.py:137] Failed to get job status.
(sky-b7d3-weichiang pid=120832) I 07-06 19:18:20 spot_state.py:189] Job failed due to unexpected controller failure.
```

```
I 07-06 19:18:16 spot_utils.py:224] SKY INFO: The return code is 255. Check the job status in 20 seconds.
I 07-06 19:18:36 spot_utils.py:234] Logs finished for job 17 (status: FAILED_CONTROLLER).

ID  NAME                RESOURCES        SUBMITTED    TOT. DURATION  JOB DURATION  #RECOVERIES  STATUS
17  sky-b7d3-weichiang  1x [V100:1]      57 mins ago  56m 44s        50m 29s       0            FAILED_CONTROLLER
``` 

I didn't encounter such error in earlier commits. Have we changed anything that could cause this issue? @Michaelvll

Hints from issue discussion:
Any useful message in `sky logs sky-spot-controller 17`?
Here is the log. There seems to be some command error for unknown reason.
```
(sky-b7d3-weichiang pid=120832) I 07-06 19:17:52 spot_utils.py:54] === Checking the job status... ===
(sky-b7d3-weichiang pid=120832) I 07-06 19:17:54 spot_utils.py:56] Job status: JobStatus.RUNNING
(sky-b7d3-weichiang pid=120832) I 07-06 19:17:54 spot_utils.py:59] ==================================
(sky-b7d3-weichiang pid=120832) I 07-06 19:18:14 spot_utils.py:54] === Checking the job status... ===
(sky-b7d3-weichiang pid=120832) E 07-06 19:18:16 subprocess_utils.py:69]
(sky-b7d3-weichiang pid=120832) E 07-06 19:18:16 controller.py:137] Unexpected error occurred: CommandError: Command python3 -u -c 'from sky.skylet import job_lib, log_lib;job_id = None if None is not None else job_lib.get_latest_job_id();job_status = job_lib.get_status(job_id);status_str = None if job_status is None else job_status.value;print("Job", job_id, status_str, flush=True)' failed with return code 255.
(sky-b7d3-weichiang pid=120832) E 07-06 19:18:16 controller.py:137] Failed to get job status.
(sky-b7d3-weichiang pid=120832) I 07-06 19:18:20 spot_state.py:189] Job failed due to unexpected controller failure.
```
 I login and run it manually. it showed `Job 17 SUCCEEDED`. Not sure if it's expected.
```
ubuntu@ip-172-31-92-239:~$ python3 -u -c 'from sky.skylet import job_lib, log_lib;job_id = None if None is not None else job_lib.get_latest_job_id();job_status = job_lib.get_status(job_id);status_str = None if job_status is None else job_status.value;print("Job", job_id, status_str, flush=True)'
Job 17 SUCCEEDED
```
The return code 255 seems like a network connection issue. We may want to see if we need some retries for job status checking.

Commit message: Fix

Patch to review:
diff --git a/sky/backends/cloud_vm_ray_backend.py b/sky/backends/cloud_vm_ray_backend.py
index 2d1e9311de0..f128875ecaf 100644
--- a/sky/backends/cloud_vm_ray_backend.py
+++ b/sky/backends/cloud_vm_ray_backend.py
@@ -2001,8 +2001,8 @@ def tail_logs(self,
                 # Refer to: https://github.com/ray-project/ray/blob/d462172be7c5779abf37609aed08af112a533e1e/python/ray/autoscaler/_private/subprocess_output_util.py#L264 # pylint: disable=line-too-long
                 stdin=subprocess.DEVNULL,
             )
-        except SystemExit as e:
-            returncode = e.code
+        except KeyboardInterrupt as e:
+            returncode = e.args[0]
         return returncode
 
     def tail_spot_logs(self,
diff --git a/sky/spot/recovery_strategy.py b/sky/spot/recovery_strategy.py
index 7c8e20d8020..495b281951f 100644
--- a/sky/spot/recovery_strategy.py
+++ b/sky/spot/recovery_strategy.py
@@ -77,7 +77,7 @@ def launch(self,
                            cluster_name=self.cluster_name,
                            detach_run=True)
                 logger.info('Spot cluster launched.')
-            except SystemExit:
+            except Exception:  # pylint: disable=broad-except
                 # If the launch fails, it will be recovered by the following
                 # code.
                 logger.info('Failed to launch the spot cluster.')
@@ -149,7 +149,7 @@ def recover(self) -> float:
             self.cluster_name)
         try:
             self.backend.cancel_jobs(handle, jobs=None)
-        except SystemExit:
+        except exceptions.CommandError:
             # Ignore the failure as the cluster can be totally stopped, and the
             # job canceling can get connection error.
             logger.info('Ignoring the job cancellation failure; the spot '
diff --git a/sky/spot/spot_utils.py b/sky/spot/spot_utils.py
index 2ceb7b501c3..244e4dd1041 100644
--- a/sky/spot/spot_utils.py
+++ b/sky/spot/spot_utils.py
@@ -11,6 +11,7 @@
 import filelock
 
 from sky import backends
+from sky import exceptions
 from sky import global_user_state
 from sky import sky_logging
 from sky.backends import backend_utils
@@ -54,7 +55,7 @@ def get_job_status(backend: 'backends.CloudVmRayBackend',
         logger.info('=== Checking the job status... ===')
         status = backend.get_job_status(handle, stream_logs=False)
         logger.info(f'Job status: {status}')
-    except SystemExit:
+    except exceptions.CommandError:
         logger.info('Failed to connect to the cluster.')
     logger.info('=' * 34)
     return status
diff --git a/tests/test_smoke.py b/tests/test_smoke.py
index 5b3d6e0f457..858325968c2 100644
--- a/tests/test_smoke.py
+++ b/tests/test_smoke.py
@@ -5,6 +5,7 @@
 import tempfile
 import time
 from typing import Dict, List, NamedTuple, Optional, Tuple
+import uuid
 
 import colorama
 import pytest
@@ -18,6 +19,7 @@
 # (username, last 4 chars of hash of hostname): for uniquefying users on
 # shared-account cloud providers.
 _smoke_test_hash = backend_utils.user_and_hostname_hash()
+test_id = str(uuid.uuid4())[-2:]
 
 
 class Test(NamedTuple):
@@ -45,7 +47,7 @@ def _get_cluster_name():
     """
     caller_func_name = inspect.stack()[1][3]
     test_name = caller_func_name.replace('_', '-')
-    return f'{test_name}-{_smoke_test_hash}'
+    return f'{test_name}-{_smoke_test_hash}-{test_id}'
 
 
 def run_one_test(test: Test) -> Tuple[int, str, str]:
@@ -497,6 +499,33 @@ def test_gcp_spot():
     run_one_test(test)
 
 
+# ---------- Testing managed spot recovery ----------
+def test_spot_recovery():
+    """Test managed spot recovery."""
+    name = _get_cluster_name()
+    region = 'us-west-2'
+    test = Test(
+        'managed-spot-recovery',
+        [
+            f'sky spot launch --region {region} -n {name} examples/managed_spot.yaml -y -d',
+            'sleep 30',
+            f's=$(sky spot status); printf "$s"; echo; echo; printf "$s" | grep {name} | head -n1 | grep "RUNNING\|STARTING"',
+            # Terminate the cluster manually.
+            f'aws ec2 terminate-instances --region {region} --instance-ids $('
+            f'aws ec2 describe-instances --region {region} '
+            f'--filters Name=tag:ray-cluster-name,Values={name}* '
+            f'--query Reservations[].Instances[].InstanceId '
+            '--output text)',
+            'sleep 30',
+            f's=$(sky spot status); printf "$s"; echo; echo; printf "$s" | grep {name} | head -n1 | grep "RECOVERING\|STARTING"',
+            'sleep 30',
+            f's=$(sky spot status); printf "$s"; echo; echo; printf "$s" | grep {name} | head -n1 | grep "RUNNING\|STARTING"',
+        ],
+        f'sky spot cancel -y -n {name}',
+    )
+    run_one_test(test)
+
+
 # ---------- Testing storage for managed spot ----------
 def test_spot_storage():
     """Test storage with managed spot"""


Code review comments:
diff_hunk: @@ -77,7 +77,7 @@ def launch(self,
                            cluster_name=self.cluster_name,
                            detach_run=True)
                 logger.info('Spot cluster launched.')
-            except SystemExit:
+            except Exception:  # pylint: disable=broad-except
                 # If the launch fails, it will be recovered by the following
                 # code.
                 logger.info('Failed to launch the spot cluster.')
path: sky/spot/recovery_strategy.py
text: @user1:
do we want to show error message from the exception?

@author:
Good point! Let me add it.