Repository: skypilot-org/skypilot
Pull request: [Spot] Let the controller aware of the failed setup and fail early
Closes #1397 

When the spot job fails during setup, the master branch will retry 3 times (or forever if `--retry-until-up` is specified), and set the spot job to `FAILED_NO_RESOURCES`, which is not desired. 

This PR makes the controller aware of the failure of the setup early, and set the spot job to `FAILED`.

A side benefit: the setup logs will go to `sky spot logs` as well, so the user don't have to use `sky logs sky-spot-controller` too see the output of their job's setup.

Tested:
- [x] `sky spot launch -n test-setup-failure ./tests/test_yamls/failed_setup.yaml`
- [x] `pytest tests/test_smoke.py` (except for #1428)

Problem statement:
[Spot] Error in the setup will keep the spot cluster retrying forever
From the user:
>  The specific use case was that there was actually some errors in the setup script I had to fix, and I had to use the sky logs to see what’s going on. Otherwise seems like the spot job just keeps restarting itself in a dead loop with no feedback to the user.

TODO:
- [ ] Stream the logs for setup to `sky spot logs` as well.
- [ ] Detect the failure of the setup and fail fast.

Hints from issue discussion:

Commit message: Add failed setup status for spot

Patch to review:
diff --git a/sky/skylet/job_lib.py b/sky/skylet/job_lib.py
index 62909440621..0d16437260e 100644
--- a/sky/skylet/job_lib.py
+++ b/sky/skylet/job_lib.py
@@ -291,7 +291,7 @@ def get_latest_job_id() -> Optional[int]:
 
 
 def get_job_time_payload(job_id: int, is_end: bool) -> Optional[int]:
-    field = 'end_at' if is_end else 'start_at'
+    field = 'end_at' if is_end else 'submitted_at'
     rows = _CURSOR.execute(f'SELECT {field} FROM jobs WHERE job_id=(?)',
                            (job_id,))
     for (timestamp,) in rows:
diff --git a/sky/spot/controller.py b/sky/spot/controller.py
index 969069f4b13..1006433bda2 100644
--- a/sky/spot/controller.py
+++ b/sky/spot/controller.py
@@ -120,7 +120,9 @@ def _run(self):
                 if job_status is not None and not job_status.is_terminal():
                     # The multi-node job is still running, continue monitoring.
                     continue
-                elif job_status == job_lib.JobStatus.FAILED:
+                elif job_status in [
+                        job_lib.JobStatus.FAILED, job_lib.JobStatus.FAILED_SETUP
+                ]:
                     # The user code has probably crashed, fail immediately.
                     end_time = spot_utils.get_job_timestamp(self._backend,
                                                             self._cluster_name,
@@ -132,10 +134,12 @@ def _run(self):
                                             None,
                                             spot_job_id=self._job_id)
                     logger.info(f'\n== End of logs (ID: {self._job_id}) ==')
-                    spot_state.set_failed(
-                        self._job_id,
-                        failure_type=spot_state.SpotStatus.FAILED,
-                        end_time=end_time)
+                    status_to_set = spot_state.SpotStatus.FAILED
+                    if job_status == job_lib.JobStatus.FAILED_SETUP:
+                        status_to_set = spot_state.SpotStatus.FAILED_SETUP
+                    spot_state.set_failed(self._job_id,
+                                          failure_type=status_to_set,
+                                          end_time=end_time)
                     break
                 # Although the cluster is healthy, we fail to access the
                 # job status. Try to recover the job (will not restart the
diff --git a/sky/spot/recovery_strategy.py b/sky/spot/recovery_strategy.py
index 166c8df969c..146ccdaf00b 100644
--- a/sky/spot/recovery_strategy.py
+++ b/sky/spot/recovery_strategy.py
@@ -152,8 +152,11 @@ def _launch(self, max_retry=3, raise_on_failure=True) -> Optional[float]:
             exception = None
             try:
                 usage_lib.messages.usage.set_internal()
+                # Detach setup, so that the setup failure can be detected
+                # by the controller process (job_status -> FAILED_SETUP).
                 sky.launch(self.dag,
                            cluster_name=self.cluster_name,
+                           detach_setup=True,
                            detach_run=True,
                            _is_launched_by_spot_controller=True)
                 logger.info('Spot cluster launched.')
@@ -223,7 +226,7 @@ def _launch(self, max_retry=3, raise_on_failure=True) -> Optional[float]:
                     continue
 
                 # Check the job status until it is not in initialized status
-                if status is not None and job_lib.JobStatus.PENDING < status:
+                if status is not None and job_lib.JobStatus.INIT < status:
                     try:
                         launch_time = spot_utils.get_job_timestamp(
                             self.backend, self.cluster_name, get_end_time=False)
diff --git a/sky/spot/spot_state.py b/sky/spot/spot_state.py
index 42b8ff3b01d..0e515c4f2bd 100644
--- a/sky/spot/spot_state.py
+++ b/sky/spot/spot_state.py
@@ -60,6 +60,7 @@ class SpotStatus(enum.Enum):
     # Terminal statuses
     SUCCEEDED = 'SUCCEEDED'
     FAILED = 'FAILED'
+    FAILED_SETUP = 'FAILED_SETUP'
     FAILED_NO_RESOURCE = 'FAILED_NO_RESOURCE'
     FAILED_CONTROLLER = 'FAILED_CONTROLLER'
     CANCELLED = 'CANCELLED'
@@ -83,7 +84,10 @@ def terminal_statuses(cls) -> List['SpotStatus']:
 
     @classmethod
     def failure_statuses(cls) -> List['SpotStatus']:
-        return [cls.FAILED, cls.FAILED_NO_RESOURCE, cls.FAILED_CONTROLLER]
+        return [
+            cls.FAILED, cls.FAILED_SETUP, cls.FAILED_NO_RESOURCE,
+            cls.FAILED_CONTROLLER
+        ]
 
 
 _SPOT_STATUS_TO_COLOR = {
@@ -94,6 +98,7 @@ def failure_statuses(cls) -> List['SpotStatus']:
     SpotStatus.RECOVERING: colorama.Fore.CYAN,
     SpotStatus.SUCCEEDED: colorama.Fore.GREEN,
     SpotStatus.FAILED: colorama.Fore.RED,
+    SpotStatus.FAILED_SETUP: colorama.Fore.RED,
     SpotStatus.FAILED_NO_RESOURCE: colorama.Fore.RED,
     SpotStatus.FAILED_CONTROLLER: colorama.Fore.RED,
     SpotStatus.CANCELLED: colorama.Fore.YELLOW,
diff --git a/sky/spot/spot_utils.py b/sky/spot/spot_utils.py
index 3fc0f6dd834..51c0ee7e31b 100644
--- a/sky/spot/spot_utils.py
+++ b/sky/spot/spot_utils.py
@@ -42,6 +42,7 @@
                                '{status_str}.[/] It may take a few minutes.')
 _JOB_CANCELLED_MESSAGE = ('[bold cyan]Waiting for the job status to be updated.'
                           '[/] It may take a minute.')
+_FINAL_SPOT_STATUS_WAIT_TIMEOUT_SECONDS = 10
 
 
 class UserSignal(enum.Enum):
@@ -317,10 +318,13 @@ def stream_logs_by_id(job_id: int, follow: bool = True) -> str:
     # The spot_status may not be in terminal status yet, since the controllerhas
     # not updated the spot state yet. We wait for a while, until the spot state
     # is updated.
+    wait_seconds = 0
     spot_status = spot_state.get_status(job_id)
     assert spot_status is not None, job_id
-    while not spot_status.is_terminal() and follow:
+    while (not spot_status.is_terminal() and follow and
+           wait_seconds < _FINAL_SPOT_STATUS_WAIT_TIMEOUT_SECONDS):
         time.sleep(1)
+        wait_seconds += 1
         spot_status = spot_state.get_status(job_id)
         assert spot_status is not None, job_id
 
diff --git a/tests/test_smoke.py b/tests/test_smoke.py
index 5b805deb0cc..debe00f6f5b 100644
--- a/tests/test_smoke.py
+++ b/tests/test_smoke.py
@@ -1103,8 +1103,6 @@ def test_use_spot(generic_cloud: str):
 def test_spot(generic_cloud: str):
     """Test the spot yaml."""
     name = _get_cluster_name()
-    cancel_command = (
-        f'sky spot cancel -y -n {name}-1; sky spot cancel -y -n {name}-2')
     test = Test(
         'managed-spot',
         [
@@ -1119,7 +1117,27 @@ def test_spot(generic_cloud: str):
             'sleep 200',
             f'{_SPOT_QUEUE_WAIT}| grep {name}-2 | head -n1 | grep "RUNNING\|SUCCEEDED"',
         ],
-        cancel_command,
+        f'sky spot cancel -y -n {name}-1; sky spot cancel -y -n {name}-2',
+        # Increase timeout since sky spot queue -r can be blocked by other spot tests.
+        timeout=20 * 60,
+    )
+    run_one_test(test)
+
+
+def test_spot_failed_setup():
+    """Test managed spot job with failed setup."""
+    name = _get_cluster_name()
+    test = Test(
+        'spot-failed-setup',
+        [
+            f'sky spot launch -n {name} -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'sky spot cancel -y -n {name}',
+        # Increase timeout since sky spot queue -r can be blocked by other spot tests.
+        timeout=20 * 60,
     )
     run_one_test(test)
 
@@ -1411,6 +1429,8 @@ def test_spot_storage(generic_cloud: str):
                 f'[ $(aws s3api list-buckets --query "Buckets[?contains(Name, \'{storage_name}\')].Name" --output text | wc -l) -eq 0 ]'
             ],
             f'sky spot cancel -y -n {name}',
+            # Increase timeout since sky spot queue -r can be blocked by other spot tests.
+            timeout=20 * 60,
         )
         run_one_test(test)
 
@@ -1431,40 +1451,44 @@ def test_spot_tpu():
             f'{_SPOT_QUEUE_WAIT}| grep {name} | head -n1 | grep "RUNNING\|SUCCEEDED"',
         ],
         f'sky spot cancel -y -n {name}',
+        # Increase timeout since sky spot queue -r can be blocked by other spot tests.
+        timeout=20 * 60,
     )
     run_one_test(test)
 
 
-# ---------- Testing env ----------
-def test_inline_env(generic_cloud: str):
-    """Test env"""
+# ---------- Testing env for spot ----------
+@pytest.mark.managed_spot
+def test_spot_inline_env(generic_cloud: str):
+    """Test spot env"""
     name = _get_cluster_name()
     test = Test(
-        'test-inline-env',
+        'test-spot-inline-env',
         [
-            f'sky launch -c {name} -y --cloud {generic_cloud} --env TEST_ENV="hello world" -- "([[ ! -z \\"\$TEST_ENV\\" ]] && [[ ! -z \\"\$SKYPILOT_NODE_IPS\\" ]] && [[ ! -z \\"\$SKYPILOT_NODE_RANK\\" ]]) || exit 1"',
-            f'sky logs {name} 1 --status',
-            f'sky exec {name} --env TEST_ENV2="success" "([[ ! -z \\"\$TEST_ENV2\\" ]] && [[ ! -z \\"\$SKYPILOT_NODE_IPS\\" ]] && [[ ! -z \\"\$SKYPILOT_NODE_RANK\\" ]]) || exit 1"',
-            f'sky logs {name} 2 --status',
+            f'sky spot launch -n {name} -y --cloud {generic_cloud} --env TEST_ENV="hello world" -- "([[ ! -z \\"\$TEST_ENV\\" ]] && [[ ! -z \\"\$SKYPILOT_NODE_IPS\\" ]] && [[ ! -z \\"\$SKYPILOT_NODE_RANK\\" ]]) || exit 1"',
+            'sleep 20',
+            f'{_SPOT_QUEUE_WAIT} | grep {name} | grep SUCCEEDED',
         ],
-        f'sky down -y {name}',
+        f'sky spot cancel -y -n {name}',
+        # Increase timeout since sky spot queue -r can be blocked by other spot tests.
+        timeout=20 * 60,
     )
     run_one_test(test)
 
 
-# ---------- Testing env for spot ----------
-@pytest.mark.managed_spot
-def test_inline_spot_env(generic_cloud: str):
+# ---------- Testing env ----------
+def test_inline_env(generic_cloud: str):
     """Test env"""
     name = _get_cluster_name()
     test = Test(
-        'test-inline-spot-env',
+        'test-inline-env',
         [
-            f'sky spot launch -n {name} -y --cloud {generic_cloud} --env TEST_ENV="hello world" -- "([[ ! -z \\"\$TEST_ENV\\" ]] && [[ ! -z \\"\$SKYPILOT_NODE_IPS\\" ]] && [[ ! -z \\"\$SKYPILOT_NODE_RANK\\" ]]) || exit 1"',
-            'sleep 20',
-            f'{_SPOT_QUEUE_WAIT} | grep {name} | grep SUCCEEDED',
+            f'sky launch -c {name} -y --cloud {generic_cloud} --env TEST_ENV="hello world" -- "([[ ! -z \\"\$TEST_ENV\\" ]] && [[ ! -z \\"\$SKYPILOT_NODE_IPS\\" ]] && [[ ! -z \\"\$SKYPILOT_NODE_RANK\\" ]]) || exit 1"',
+            f'sky logs {name} 1 --status',
+            f'sky exec {name} --env TEST_ENV2="success" "([[ ! -z \\"\$TEST_ENV2\\" ]] && [[ ! -z \\"\$SKYPILOT_NODE_IPS\\" ]] && [[ ! -z \\"\$SKYPILOT_NODE_RANK\\" ]]) || exit 1"',
+            f'sky logs {name} 2 --status',
         ],
-        f'sky spot cancel -y -n {name}',
+        f'sky down -y {name}',
     )
     run_one_test(test)
 
diff --git a/tests/test_yamls/failed_setup.yaml b/tests/test_yamls/failed_setup.yaml
new file mode 100644
index 00000000000..f6131e7c204
--- /dev/null
+++ b/tests/test_yamls/failed_setup.yaml
@@ -0,0 +1,6 @@
+setup: |
+  echo failed setup
+  exit 1
+
+run: |
+  echo run


Code review comments:
diff_hunk: @@ -223,7 +226,7 @@ def _launch(self, max_retry=3, raise_on_failure=True) -> Optional[float]:
                     continue
 
                 # Check the job status until it is not in initialized status
-                if status is not None and job_lib.JobStatus.PENDING < status:
+                if status is not None and job_lib.JobStatus.INIT < status:
path: sky/spot/recovery_strategy.py
text: @user1:
```suggestion
                if status is not None and status > job_lib.JobStatus.INIT:
```

diff_hunk: @@ -291,7 +291,7 @@ def get_latest_job_id() -> Optional[int]:
 
 
 def get_job_time_payload(job_id: int, is_end: bool) -> Optional[int]:
path: sky/skylet/job_lib.py
text: @user1:
Can we rename this? `get_job_submitted_or_ended_timestamp_payload`? "job time" sounds like a duration, but here we want one of the two timestamps.

Need to update https://github.com/skypilot-org/skypilot/blob/8a80f5be49d951d1e6334f2cd1e6595f65e4e413/sky/spot/spot_utils.py#L123 comment to say submitted_at?  

To keep this interface consistent with that func, is_end -> get_end_time?

@author:
Good point! Updated. Thanks!

diff_hunk: @@ -291,7 +291,7 @@ def get_latest_job_id() -> Optional[int]:
 
 
 def get_job_time_payload(job_id: int, is_end: bool) -> Optional[int]:
-    field = 'end_at' if is_end else 'start_at'
+    field = 'end_at' if is_end else 'submitted_at'
path: sky/skylet/job_lib.py
text: @user1:
Should we be concerned about this change?

Many places encode the return value as start_at, which has downstream usage/assumptions:
- https://github.com/skypilot-org/skypilot/blob/8a80f5be49d951d1e6334f2cd1e6595f65e4e413/sky/spot/controller.py#L67
- etc

@user1:
Also, confused why we need to do this change. Let's say we submit 1000 spot jobs, thus making many of them pending. Would this change somehow start counting those pending jobs' duration?

@author:
> Should we be concerned about this change?
> Many places encode the return value as start_at, which has downstream usage/assumptions:
> https://github.com/skypilot-org/skypilot/blob/8a80f5be49d951d1e6334f2cd1e6595f65e4e413/sky/spot/controller.py#L67

Thanks for catching this! I changed the variable name to `launched_at` instead, but still keep the concept for spot job to be `start_time`, as we now consider the setup to happen after the the spot job starts.

> Also, confused why we need to do this change. Let's say we submit 1000 spot jobs, thus making many of them pending. Would this change somehow start counting those pending jobs' duration?

The pending period for the spot job does not count for the jobs' duration, because:
1. Each spot cluster will be dedicated for a single spot job, which means the job will not be in JobStatus.PENDING mode when the spot cluster is healthy.
2. 1000 spot jobs will be pending on the controller, i.e. in the job queue of the controller. The submitted_at/start_at are only retrieved from the spot cluster, after the spot job is actually submitted to a spot cluster. The only change to the job duration is from the duration for the `run` section to the duration of all `setup` and `run` section.