Fix spot terminal status (skypilot-org#1624)

* Fix spot terminal status

* fix smoke test

* longer waiting time

* longer

* fix FAILED_SETUP type

* format

* fix comment

Author: Michaelvll

sky/spot/spot_state.py

@@ -127,7 +127,7 @@ def colored_str(self):
     @classmethod
     def terminal_statuses(cls) -> List['SpotStatus']:
         return [
-            cls.SUCCEEDED, cls.FAILED, cls.FAILED_NO_RESOURCE,
+            cls.SUCCEEDED, cls.FAILED, cls.FAILED_SETUP, cls.FAILED_NO_RESOURCE,
             cls.FAILED_CONTROLLER, cls.CANCELLED
         ]
 
@@ -264,8 +264,9 @@ def set_failed(job_id: int,
         WHERE job_id=(?) AND end_at IS null""",
         (*list(fields_to_set.values()), job_id))
     _CONN.commit()
-    if failure_type == SpotStatus.FAILED:
-        logger.info('Job failed due to user code.')
+    if failure_type in [SpotStatus.FAILED, SpotStatus.FAILED_SETUP]:
+        logger.info(
+            f'Job failed due to user code (status: {failure_type.value}).')
     elif failure_type == SpotStatus.FAILED_NO_RESOURCE:
         logger.info('Job failed due to failing to find available resources '
                     'after retries.')

sky/spot/spot_utils.py

@@ -47,7 +47,7 @@
 # state, after the job finished. This is a safeguard to avoid the case where
 # the spot job status fails to be updated and keep the `sky spot logs` blocking
 # for a long time.
-_FINAL_SPOT_STATUS_WAIT_TIMEOUT_SECONDS = 10
+_FINAL_SPOT_STATUS_WAIT_TIMEOUT_SECONDS = 20
 
 
 class UserSignal(enum.Enum):
@@ -62,7 +62,7 @@ def get_job_status(backend: 'backends.CloudVmRayBackend',
                    cluster_name: str) -> Optional['job_lib.JobStatus']:
     """Check the status of the job running on the spot cluster.
 
-    It can be None, INIT, RUNNING, SUCCEEDED, FAILED or CANCELLED.
+    It can be None, INIT, RUNNING, SUCCEEDED, FAILED, FAILED_SETUP or CANCELLED.
     """
     handle = global_user_state.get_handle_from_cluster_name(cluster_name)
     status = None

tests/test_smoke.py

@@ -1134,7 +1134,7 @@ def test_spot_failed_setup(generic_cloud: str):
             f'sky spot launch -n {name} --cloud {generic_cloud} -y -d tests/test_yamls/failed_setup.yaml',
             'sleep 200',
             # Make sure the job failed quickly.
-            f'{_SPOT_QUEUE_WAIT} | grep {name} | head -n1 | grep "FAILED"',
+            f'{_SPOT_QUEUE_WAIT} | grep {name} | head -n1 | grep "FAILED_SETUP"',
         ],
         f'sky spot cancel -y -n {name}',
         # Increase timeout since sky spot queue -r can be blocked by other spot tests.