Repository: skypilot-org/skypilot
Pull request: [Spot] Fix spot recovery for multi node
Closes #1410 

Tested:
- [x] `tests/run_smoke_tests.sh test_spot_recovery_multi_node`

Problem statement:
[Spot] Controller not recover multi-node spot job when preempted
When the worker node is preempted, the controller will not recover the cluster, if the job does not fail immediately with non-zero returncode after the number of nodes becomes less than expected.

Hints from issue discussion:

Commit message: Add test

Patch to review:
diff --git a/sky/spot/controller.py b/sky/spot/controller.py
index 6b7a452fbfc..0487da3a1fd 100644
--- a/sky/spot/controller.py
+++ b/sky/spot/controller.py
@@ -86,8 +86,21 @@ def _run(self):
                                                    self._cluster_name)
 
             if job_status is not None and not job_status.is_terminal():
-                # The job is healthy, continue to monitor the job status.
-                continue
+                need_recovery = False
+                if self._task.num_nodes > 1:
+                    # Check the cluster status for multi-node jobs, since the
+                    # job may not FAILED immediately when part of the nodes are
+                    # preempted.
+                    (cluster_status,
+                     handle) = backend_utils.refresh_cluster_status_handle(
+                         self._cluster_name, force_refresh=True)
+                    if cluster_status != global_user_state.ClusterStatus.UP:
+                        # recover the cluster if it is not up.
+                        need_recovery = True
+                if not need_recovery:
+                    # The job and cluster are healthy, continue to monitor the
+                    # job status.
+                    continue
 
             if job_status == job_lib.JobStatus.SUCCEEDED:
                 end_time = spot_utils.get_job_timestamp(self.backend,
@@ -97,12 +110,9 @@ def _run(self):
                 spot_state.set_succeeded(self._job_id, end_time=end_time)
                 break
 
-            assert (job_status is None or
-                    job_status == job_lib.JobStatus.FAILED), (
-                        f'The job should not be {job_status.value}.')
             if job_status == job_lib.JobStatus.FAILED:
-                # Check the status of the spot cluster. It can be STOPPED or UP,
-                # where STOPPED means partially down.
+                # Check the status of the spot cluster. If it is not UP,
+                # the cluster is preempted
                 (cluster_status,
                  handle) = backend_utils.refresh_cluster_status_handle(
                      self._cluster_name, force_refresh=True)
diff --git a/sky/spot/spot_utils.py b/sky/spot/spot_utils.py
index a67ab4f8d8f..53b5dbba87d 100644
--- a/sky/spot/spot_utils.py
+++ b/sky/spot/spot_utils.py
@@ -260,12 +260,17 @@ def stream_logs_by_id(job_id: int, follow: bool = True) -> str:
             job_statuses = backend.get_job_status(handle, stream_logs=False)
             job_status = list(job_statuses.values())[0]
             assert job_status is not None, 'No job found.'
-            logger.info(f'Logs finished for job {job_id} '
-                        f'(status: {job_status.value}).')
-            break
-        logger.info(
-            f'INFO: (Log streaming) Got return code {returncode}. Retrying '
-            f'in {JOB_STATUS_CHECK_GAP_SECONDS} seconds.')
+            if job_status != job_lib.JobStatus.CANCELLED:
+                logger.info(f'Logs finished for job {job_id} '
+                            f'(status: {job_status.value}).')
+                break
+            logger.info('INFO: (Log streaming) The job is cancelled. Waiting '
+                        'for the controller process to update the status in '
+                        f'{JOB_STATUS_CHECK_GAP_SECONDS} seconds.')
+        else:
+            logger.info(
+                f'INFO: (Log streaming) Got return code {returncode}. Retrying '
+                f'in {JOB_STATUS_CHECK_GAP_SECONDS} seconds.')
         # If the tailing fails, it is likely that the cluster fails, so we wait
         # a while to make sure the spot state is updated by the controller, and
         # check the spot queue again.
diff --git a/sky/templates/aws-ray.yml.j2 b/sky/templates/aws-ray.yml.j2
index de0bbdac60b..1ba5cc3ebaf 100644
--- a/sky/templates/aws-ray.yml.j2
+++ b/sky/templates/aws-ray.yml.j2
@@ -132,12 +132,12 @@ head_start_ray_commands:
   # Line "which prlimit ..": increase the limit of the number of open files for the raylet process, as the `ulimit` may not take effect at this point, because it requires
   # all the sessions to be reloaded. This is a workaround.
   - ((ps aux | grep -v nohup | grep -v grep | grep -q -- "python3 -m sky.skylet.skylet") || nohup python3 -m sky.skylet.skylet >> ~/.sky/skylet.log 2>&1 &);
-    ray stop; ray start --disable-usage-stats --head --port=6379 --object-manager-port=8076 --autoscaling-config=~/ray_bootstrap_config.yaml {{"--resources='%s'" % custom_resources if custom_resources}} || exit 1;
+    ray stop; RAY_SCHEDULER_EVENTS=0 ray start --disable-usage-stats --head --port=6379 --object-manager-port=8076 --autoscaling-config=~/ray_bootstrap_config.yaml {{"--resources='%s'" % custom_resources if custom_resources}} || exit 1;
     which prlimit && for id in $(pgrep -f raylet/raylet); do sudo prlimit --nofile=1048576:1048576 --pid=$id || true; done;
 
 {%- if num_nodes > 1 %}
 worker_start_ray_commands:
-  - ray stop; ray start --disable-usage-stats --address=$RAY_HEAD_IP:6379 --object-manager-port=8076 {{"--resources='%s'" % custom_resources if custom_resources}} || exit 1;
+  - ray stop; RAY_SCHEDULER_EVENTS=0 ray start --disable-usage-stats --address=$RAY_HEAD_IP:6379 --object-manager-port=8076 {{"--resources='%s'" % custom_resources if custom_resources}} || exit 1;
     which prlimit && for id in $(pgrep -f raylet/raylet); do sudo prlimit --nofile=1048576:1048576 --pid=$id || true; done;
 {%- else %}
 worker_start_ray_commands: []
diff --git a/sky/templates/azure-ray.yml.j2 b/sky/templates/azure-ray.yml.j2
index 5d98ef90c0b..ab3de7405c4 100644
--- a/sky/templates/azure-ray.yml.j2
+++ b/sky/templates/azure-ray.yml.j2
@@ -128,12 +128,12 @@ head_start_ray_commands:
   # Start skylet daemon. (Should not place it in the head_setup_commands, otherwise it will run before skypilot is installed.)
   # NOTE: --disable-usage-stats in `ray start` saves 10 seconds of idle wait.
   - ((ps aux | grep -v nohup | grep -v grep | grep -q -- "python3 -m sky.skylet.skylet") || nohup python3 -m sky.skylet.skylet >> ~/.sky/skylet.log 2>&1 &);
-    ray stop; ray start --disable-usage-stats --head --port=6379 --object-manager-port=8076 --autoscaling-config=~/ray_bootstrap_config.yaml {{"--resources='%s'" % custom_resources if custom_resources}} || exit 1;
+    ray stop; RAY_SCHEDULER_EVENTS=0 ray start --disable-usage-stats --head --port=6379 --object-manager-port=8076 --autoscaling-config=~/ray_bootstrap_config.yaml {{"--resources='%s'" % custom_resources if custom_resources}} || exit 1;
     which prlimit && for id in $(pgrep -f raylet/raylet); do sudo prlimit --nofile=1048576:1048576 --pid=$id || true; done;
 
 {%- if num_nodes > 1 %}
 worker_start_ray_commands:
-  - ray stop; ray start --disable-usage-stats --address=$RAY_HEAD_IP:6379 --object-manager-port=8076 {{"--resources='%s'" % custom_resources if custom_resources}} || exit 1;
+  - ray stop; RAY_SCHEDULER_EVENTS=0 ray start --disable-usage-stats --address=$RAY_HEAD_IP:6379 --object-manager-port=8076 {{"--resources='%s'" % custom_resources if custom_resources}} || exit 1;
     which prlimit && for id in $(pgrep -f raylet/raylet); do sudo prlimit --nofile=1048576:1048576 --pid=$id || true; done;
 {%- else %}
 worker_start_ray_commands: []
diff --git a/sky/templates/gcp-ray.yml.j2 b/sky/templates/gcp-ray.yml.j2
index c836cbc2f4e..620e2f008cb 100644
--- a/sky/templates/gcp-ray.yml.j2
+++ b/sky/templates/gcp-ray.yml.j2
@@ -172,14 +172,14 @@ head_start_ray_commands:
   # all the sessions to be reloaded. This is a workaround.
   - ((ps aux | grep -v nohup | grep -v grep | grep -q -- "python3 -m sky.skylet.skylet") || nohup python3 -m sky.skylet.skylet >> ~/.sky/skylet.log 2>&1 &);
     export SKYPILOT_NUM_GPUS=0 && which nvidia-smi > /dev/null && SKYPILOT_NUM_GPUS=$(nvidia-smi --query-gpu=index,name --format=csv,noheader | wc -l);
-    ray stop; ray start --disable-usage-stats --head --port=6379 --object-manager-port=8076 --autoscaling-config=~/ray_bootstrap_config.yaml {{"--resources='%s'" % custom_resources if custom_resources}} --num-gpus=$SKYPILOT_NUM_GPUS || exit 1;
+    ray stop; RAY_SCHEDULER_EVENTS=0 ray start --disable-usage-stats --head --port=6379 --object-manager-port=8076 --autoscaling-config=~/ray_bootstrap_config.yaml {{"--resources='%s'" % custom_resources if custom_resources}} --num-gpus=$SKYPILOT_NUM_GPUS || exit 1;
     which prlimit && for id in $(pgrep -f raylet/raylet); do sudo prlimit --nofile=1048576:1048576 --pid=$id || true; done;
 
 # Worker commands are needed for TPU VM Pods
 {%- if num_nodes > 1 or tpu_vm %}
 worker_start_ray_commands:
   - SKYPILOT_NUM_GPUS=0 && which nvidia-smi > /dev/null && SKYPILOT_NUM_GPUS=$(nvidia-smi --query-gpu=index,name --format=csv,noheader | wc -l);
-    ray stop; ray start --disable-usage-stats --address=$RAY_HEAD_IP:6379 --object-manager-port=8076 {{"--resources='%s'" % custom_resources if custom_resources}} --num-gpus=$SKYPILOT_NUM_GPUS || exit 1;
+    ray stop; RAY_SCHEDULER_EVENTS=0 ray start --disable-usage-stats --address=$RAY_HEAD_IP:6379 --object-manager-port=8076 {{"--resources='%s'" % custom_resources if custom_resources}} --num-gpus=$SKYPILOT_NUM_GPUS || exit 1;
     which prlimit && for id in $(pgrep -f raylet/raylet); do sudo prlimit --nofile=1048576:1048576 --pid=$id || true; done;
 {%- else %}
 worker_start_ray_commands: []
diff --git a/tests/test_smoke.py b/tests/test_smoke.py
index 3827d4fe6ef..2c925358066 100644
--- a/tests/test_smoke.py
+++ b/tests/test_smoke.py
@@ -807,7 +807,7 @@ def test_spot_recovery():
             f'sky spot launch --cloud aws --region {region} -n {name} "echo SKYPILOT_JOB_ID: \$SKYPILOT_JOB_ID; sleep 1000"  -y -d',
             'sleep 300',
             f's=$(sky spot queue); printf "$s"; echo; echo; printf "$s" | grep {name} | head -n1 | grep "RUNNING"',
-            f'RUN_ID=$(sky spot logs -n {name} --no-follow | grep SKYPILOT_JOB_ID | cut -d: -f2); echo $RUN_ID | tee /tmp/{name}-run-id',
+            f'RUN_ID=$(sky spot logs -n {name} --no-follow | grep SKYPILOT_JOB_ID | cut -d: -f2); echo "$RUN_ID" | tee /tmp/{name}-run-id',
             # Terminate the cluster manually.
             (f'aws ec2 terminate-instances --region {region} --instance-ids $('
              f'aws ec2 describe-instances --region {region} '
@@ -825,6 +825,35 @@ def test_spot_recovery():
     run_one_test(test)
 
 
+def test_spot_recovery_multi_node():
+    """Test managed spot recovery."""
+    name = _get_cluster_name()
+    region = 'us-west-2'
+    test = Test(
+        'managed-spot-recovery',
+        [
+            f'sky spot launch --cloud aws --region {region} -n {name} --num-nodes 2 "echo SKYPILOT_JOB_ID: \$SKYPILOT_JOB_ID; sleep 1000"  -y -d',
+            'sleep 360',
+            f's=$(sky spot queue); printf "$s"; echo; echo; printf "$s" | grep {name} | head -n1 | grep "RUNNING"',
+            f'RUN_ID=$(sky spot logs -n {name} --no-follow | grep SKYPILOT_JOB_ID | cut -d: -f2); echo "$RUN_ID" | tee /tmp/{name}-run-id',
+            # 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}* '
+             'Name=tag:ray-node-type,Values=worker '
+             f'--query Reservations[].Instances[].InstanceId '
+             '--output text)'),
+            'sleep 50',
+            f's=$(sky spot queue); printf "$s"; echo; echo; printf "$s" | grep {name} | head -n1 | grep "RECOVERING"',
+            'sleep 360',
+            f's=$(sky spot queue); printf "$s"; echo; echo; printf "$s" | grep {name} | head -n1 | grep "RUNNING"',
+            f'RUN_ID=$(cat /tmp/{name}-run-id); echo $RUN_ID; sky spot logs -n {name} --no-follow | grep SKYPILOT_JOB_ID | cut -d: -f2 | grep "$RUN_ID"',
+        ],
+        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: @@ -97,12 +110,9 @@ def _run(self):
                 spot_state.set_succeeded(self._job_id, end_time=end_time)
                 break
 
-            assert (job_status is None or
path: sky/spot/controller.py
text: @user1:
ok to remove this assert?

diff_hunk: @@ -260,12 +260,17 @@ def stream_logs_by_id(job_id: int, follow: bool = True) -> str:
             job_statuses = backend.get_job_status(handle, stream_logs=False)
             job_status = list(job_statuses.values())[0]
             assert job_status is not None, 'No job found.'
-            logger.info(f'Logs finished for job {job_id} '
-                        f'(status: {job_status.value}).')
-            break
-        logger.info(
-            f'INFO: (Log streaming) Got return code {returncode}. Retrying '
-            f'in {JOB_STATUS_CHECK_GAP_SECONDS} seconds.')
+            if job_status != job_lib.JobStatus.CANCELLED:
+                logger.info(f'Logs finished for job {job_id} '
+                            f'(status: {job_status.value}).')
+                break
+            logger.info('INFO: (Log streaming) The job is cancelled. Waiting '
+                        'for the controller process to update the status in '
+                        f'{JOB_STATUS_CHECK_GAP_SECONDS} seconds.')
path: sky/spot/spot_utils.py
text: @user1:
Nit: to me (a user) this seems a bit hard to understand (and does the user need to know this?).

@author:
Updated it a bit. Since we cannot distinguish the job being cancelled by controller due to partially preemption and actual user's cancellation before the spot status being updated, we have to wait a while until the spot status is updated.

diff_hunk: @@ -86,8 +86,21 @@ def _run(self):
                                                    self._cluster_name)
 
             if job_status is not None and not job_status.is_terminal():
-                # The job is healthy, continue to monitor the job status.
-                continue
+                need_recovery = False
+                if self._task.num_nodes > 1:
+                    # Check the cluster status for multi-node jobs, since the
+                    # job may not FAILED immediately when part of the nodes are
+                    # preempted.
path: sky/spot/controller.py
text: @user1:
```suggestion
                    # job may not be set to FAILED immediately when only some of the nodes are
                    # preempted.
```

@author:
Done. Thanks!