Repository: skypilot-org/skypilot
Pull request: [ManagedJobs] Support failure recovery for HA controller
<!-- Describe the changes in this PR -->

Support recovery logic for managed jobs.

Fixes #3970 

<!-- Describe the tests ran -->
<!-- Unit tests (tests/test_*.py) are part of GitHub CI; below are tests that launch on the cloud. -->

Tested (run the relevant ones):

- [ ] Code formatting: install pre-commit (auto-check on commit) or `bash format.sh`
- [ ] Any manual or new tests for this PR (please specify below)
- [ ] All smoke tests: `/smoke-test` (CI) or `pytest tests/test_smoke.py` (local)
- [ ] Relevant individual tests: `/smoke-test -k test_name` (CI) or `pytest tests/test_smoke.py::test_name` (local)
- [ ] Backward compatibility: `/quicktest-core` (CI) or `pytest tests/smoke_tests/test_backward_compat.py` (local)

<!-- CI commands (/-prefixed) can only be triggered by repo members -->


Problem statement:
How to make the Controller highly available?
I have a query related to making the controller highly available for production deployment to remove the single point of failure. I have checked the doc but didn't find anything related. Is there any guide available for this?

Hints from issue discussion:
This issue is stale because it has been open 120 days with no activity. Remove stale label or comment or this will be closed in 10 days.
This issue was closed because it has been stalled for 10 days with no activity.
Reopening for jobs controller. While HA is addressed for Serve controller in #4564, jobs controller state is still stored on the controller VM/pod. If the pod/VM crashes, all jobs state is lost.

Commit message: add file

Patch to review:
diff --git a/charts/skypilot/values.yaml b/charts/skypilot/values.yaml
index 21c7438d801..b2c50b02ad5 100644
--- a/charts/skypilot/values.yaml
+++ b/charts/skypilot/values.yaml
@@ -196,6 +196,13 @@ rbac:
     - apiGroups: [ "" ]
       resources: [ "configmaps" ]
       verbs: [ "get", "patch" ]
+    # Required for high-availability controller
+    - apiGroups: ["apps"]
+      resources: ["deployments", "deployments/status"]
+      verbs: ["*"]
+    - apiGroups: [""]
+      resources: ["persistentvolumeclaims"]
+      verbs: ["*"]
   # Cluster-scoped rules for API server.
   clusterRules:
     # Required for getting node resources.
diff --git a/sky/clouds/kubernetes.py b/sky/clouds/kubernetes.py
index 5fe96e5478f..d4510280952 100644
--- a/sky/clouds/kubernetes.py
+++ b/sky/clouds/kubernetes.py
@@ -618,6 +618,8 @@ def _get_image_id(resources: 'resources_lib.Resources') -> str:
                 (constants.PERSISTENT_SETUP_SCRIPT_PATH),
             'k8s_high_availability_deployment_run_script_dir':
                 (constants.PERSISTENT_RUN_SCRIPT_DIR),
+            'k8s_high_availability_restarting_signal_file':
+                (constants.PERSISTENT_RUN_RESTARTING_SIGNAL_FILE),
             'k8s_high_availability_storage_class_name':
                 (k8s_ha_storage_class_name),
             'avoid_label_keys': avoid_label_keys,
diff --git a/sky/execution.py b/sky/execution.py
index b0e795a7c34..2b3659c8210 100644
--- a/sky/execution.py
+++ b/sky/execution.py
@@ -263,8 +263,7 @@ def _execute_dag(
     if controller is not None:
         requested_features.add(
             clouds.CloudImplementationFeatures.HOST_CONTROLLERS)
-        if controller_utils.high_availability_specified(cluster_name,
-                                                        skip_warning=False):
+        if controller_utils.high_availability_specified(cluster_name):
             requested_features.add(clouds.CloudImplementationFeatures.
                                    HIGH_AVAILABILITY_CONTROLLERS)
             # If we provision a cluster that supports high availability
diff --git a/sky/jobs/controller.py b/sky/jobs/controller.py
index 7204fd4c516..8a2e8fdaeab 100644
--- a/sky/jobs/controller.py
+++ b/sky/jobs/controller.py
@@ -152,6 +152,20 @@ def _run_one_task(self, task_id: int, task: 'sky.Task') -> bool:
         Other exceptions may be raised depending on the backend.
         """
 
+        latest_task_id, status = managed_job_state.get_latest_task_id_status(
+            self._job_id)
+        is_recovery = False
+        if (latest_task_id is not None and
+                status != managed_job_state.ManagedJobStatus.PENDING):
+            assert latest_task_id >= task_id, (latest_task_id, task_id)
+            if latest_task_id > task_id:
+                logger.info(f'Task {task_id} ({task.name}) has already '
+                            'been executed. Skipping...')
+                return True
+            if latest_task_id == task_id:
+                # Start recovery.
+                is_recovery = True
+
         callback_func = managed_job_utils.event_callback_func(
             job_id=self._job_id, task_id=task_id, task=task)
         if task.run is None:
@@ -171,42 +185,48 @@ def _run_one_task(self, task_id: int, task: 'sky.Task') -> bool:
             return True
         usage_lib.messages.usage.update_task_id(task_id)
         task_id_env_var = task.envs[constants.TASK_ID_ENV_VAR]
-        submitted_at = time.time()
-        if task_id == 0:
-            submitted_at = backend_utils.get_timestamp_from_run_timestamp(
-                self._backend.run_timestamp)
         assert task.name is not None, task
         cluster_name = managed_job_utils.generate_managed_job_cluster_name(
             task.name, self._job_id)
         self._strategy_executor = recovery_strategy.StrategyExecutor.make(
             cluster_name, self._backend, task, self._job_id, task_id)
-        managed_job_state.set_starting(
-            self._job_id,
-            task_id,
-            self._backend.run_timestamp,
-            submitted_at,
-            resources_str=backend_utils.get_task_resources_str(
-                task, is_managed_job=True),
-            specs={
-                'max_restarts_on_errors':
-                    self._strategy_executor.max_restarts_on_errors
-            },
-            callback_func=callback_func)
-        logger.info(
-            f'Submitted managed job {self._job_id} (task: {task_id}, name: '
-            f'{task.name!r}); {constants.TASK_ID_ENV_VAR}: {task_id_env_var}')
-
-        logger.info('Started monitoring.')
-
-        remote_job_submitted_at = self._strategy_executor.launch()
-        assert remote_job_submitted_at is not None, remote_job_submitted_at
+        if not is_recovery:
+            submitted_at = time.time()
+            if task_id == 0:
+                submitted_at = backend_utils.get_timestamp_from_run_timestamp(
+                    self._backend.run_timestamp)
+            managed_job_state.set_starting(
+                self._job_id,
+                task_id,
+                self._backend.run_timestamp,
+                submitted_at,
+                resources_str=backend_utils.get_task_resources_str(
+                    task, is_managed_job=True),
+                specs={
+                    'max_restarts_on_errors':
+                        self._strategy_executor.max_restarts_on_errors
+                },
+                callback_func=callback_func)
+            logger.info(f'Submitted managed job {self._job_id} '
+                        f'(task: {task_id}, name: {task.name!r}); '
+                        f'{constants.TASK_ID_ENV_VAR}: {task_id_env_var}')
+
+            logger.info('Started monitoring.')
+
+            remote_job_submitted_at = self._strategy_executor.launch()
+            assert remote_job_submitted_at is not None, remote_job_submitted_at
 
-        managed_job_state.set_started(job_id=self._job_id,
-                                      task_id=task_id,
-                                      start_time=remote_job_submitted_at,
-                                      callback_func=callback_func)
+            managed_job_state.set_started(job_id=self._job_id,
+                                          task_id=task_id,
+                                          start_time=remote_job_submitted_at,
+                                          callback_func=callback_func)
 
         while True:
+            if is_recovery:
+                last_status = managed_job_state.get_job_status_with_task_id(
+                    job_id=self._job_id, task_id=task_id)
+                if last_status is not None and last_status.is_terminal():
+                    return True
             time.sleep(managed_job_utils.JOB_STATUS_CHECK_GAP_SECONDS)
 
             # Check the network connection to avoid false alarm for job failure.
@@ -221,8 +241,12 @@ def _run_one_task(self, task_id: int, task: 'sky.Task') -> bool:
 
             # NOTE: we do not check cluster status first because race condition
             # can occur, i.e. cluster can be down during the job status check.
-            job_status = managed_job_utils.get_job_status(
-                self._backend, cluster_name)
+            try:
+                job_status = managed_job_utils.get_job_status(
+                    self._backend, cluster_name)
+            except (exceptions.FetchClusterInfoError, exceptions.CommandError):
+                logger.info('Failed to fetch the job status. Start recovery...')
+                job_status = None
 
             if job_status == job_lib.JobStatus.SUCCEEDED:
                 success_end_time = managed_job_utils.try_to_get_job_end_time(
diff --git a/sky/jobs/scheduler.py b/sky/jobs/scheduler.py
index 617c5b790b3..199c4348bf3 100644
--- a/sky/jobs/scheduler.py
+++ b/sky/jobs/scheduler.py
@@ -84,6 +84,32 @@ def _get_lock_path() -> str:
     return path
 
 
+def _start_controller(job_id: int, dag_yaml_path: str,
+                      env_file_path: str) -> None:
+    activate_python_env_cmd = (f'{constants.ACTIVATE_SKY_REMOTE_PYTHON_ENV};')
+    source_environment_cmd = (f'source {env_file_path};'
+                              if env_file_path else '')
+    run_controller_cmd = ('python -u -m sky.jobs.controller '
+                          f'{dag_yaml_path} --job-id {job_id};')
+
+    # If the command line here is changed, please also update
+    # utils._controller_process_alive. `--job-id X` should be at
+    # the end.
+    run_cmd = (f'{activate_python_env_cmd}'
+               f'{source_environment_cmd}'
+               f'{run_controller_cmd}')
+
+    logs_dir = os.path.expanduser(
+        managed_job_constants.JOBS_CONTROLLER_LOGS_DIR)
+    os.makedirs(logs_dir, exist_ok=True)
+    log_path = os.path.join(logs_dir, f'{job_id}.log')
+
+    pid = subprocess_utils.launch_new_process_tree(run_cmd, log_output=log_path)
+    state.set_job_controller_pid(job_id, pid)
+
+    logger.debug(f'Job {job_id} started with pid {pid}')
+
+
 def maybe_schedule_next_jobs() -> None:
     """Determine if any managed jobs can be scheduled, and if so, schedule them.
 
@@ -158,32 +184,9 @@ def maybe_schedule_next_jobs() -> None:
 
                     job_id = maybe_next_job['job_id']
                     dag_yaml_path = maybe_next_job['dag_yaml_path']
+                    env_file_path = maybe_next_job['env_file_path']
 
-                    activate_python_env_cmd = (
-                        f'{constants.ACTIVATE_SKY_REMOTE_PYTHON_ENV};')
-                    env_file = maybe_next_job['env_file_path']
-                    source_environment_cmd = (f'source {env_file};'
-                                              if env_file else '')
-                    run_controller_cmd = ('python -u -m sky.jobs.controller '
-                                          f'{dag_yaml_path} --job-id {job_id};')
-
-                    # If the command line here is changed, please also update
-                    # utils._controller_process_alive. `--job-id X` should be at
-                    # the end.
-                    run_cmd = (f'{activate_python_env_cmd}'
-                               f'{source_environment_cmd}'
-                               f'{run_controller_cmd}')
-
-                    logs_dir = os.path.expanduser(
-                        managed_job_constants.JOBS_CONTROLLER_LOGS_DIR)
-                    os.makedirs(logs_dir, exist_ok=True)
-                    log_path = os.path.join(logs_dir, f'{job_id}.log')
-
-                    pid = subprocess_utils.launch_new_process_tree(
-                        run_cmd, log_output=log_path)
-                    state.set_job_controller_pid(job_id, pid)
-
-                    logger.debug(f'Job {job_id} started with pid {pid}')
+                    _start_controller(job_id, dag_yaml_path, env_file_path)
 
     except filelock.Timeout:
         # If we can't get the lock, just exit. The process holding the lock
@@ -203,9 +206,14 @@ def submit_job(job_id: int, dag_yaml_path: str, env_file_path: str,
     The user hash should be set (e.g. via SKYPILOT_USER_ID) before calling this.
     """
     with filelock.FileLock(_get_lock_path()):
-        state.scheduler_set_waiting(job_id, dag_yaml_path, env_file_path,
-                                    common_utils.get_user_hash(), priority)
-    maybe_schedule_next_jobs()
+        is_recovery = state.scheduler_set_waiting(job_id, dag_yaml_path,
+                                                  env_file_path,
+                                                  common_utils.get_user_hash(),
+                                                  priority)
+    if is_recovery:
+        _start_controller(job_id, dag_yaml_path, env_file_path)
+    else:
+        maybe_schedule_next_jobs()
 
 
 @contextlib.contextmanager
diff --git a/sky/jobs/state.py b/sky/jobs/state.py
index 2a030d93bd0..da4b4e79764 100644
--- a/sky/jobs/state.py
+++ b/sky/jobs/state.py
@@ -573,6 +573,9 @@ def set_started(job_id: int, task_id: int, start_time: float,
 def set_recovering(job_id: int, task_id: int, callback_func: CallbackType):
     """Set the task to recovering state, and update the job duration."""
     logger.info('=== Recovering... ===')
+    # Originally, we force the status to be RUNNING before setting to RECOVERING
+    # After adding the HA job controller, it is possible that the jobs came from
+    # any status to recovering. So we skip the status check here.
     with db_utils.safe_cursor(_DB_PATH) as cursor:
         cursor.execute(
             """\
@@ -580,10 +583,8 @@ def set_recovering(job_id: int, task_id: int, callback_func: CallbackType):
                 status=(?), job_duration=job_duration+(?)-last_recovered_at
                 WHERE spot_job_id=(?) AND
                 task_id=(?) AND
-                status=(?) AND
                 end_at IS null""",
-            (ManagedJobStatus.RECOVERING.value, time.time(), job_id, task_id,
-             ManagedJobStatus.RUNNING.value))
+            (ManagedJobStatus.RECOVERING.value, time.time(), job_id, task_id))
         if cursor.rowcount != 1:
             raise exceptions.ManagedJobStatusError(
                 f'Failed to set the task to recovering. '
@@ -935,6 +936,17 @@ def _get_all_task_ids_statuses(
         return [(row[0], ManagedJobStatus(row[1])) for row in id_statuses]
 
 
+def get_job_status_with_task_id(job_id: int,
+                                task_id: int) -> Optional[ManagedJobStatus]:
+    with db_utils.safe_cursor(_DB_PATH) as cursor:
+        status = cursor.execute(
+            """\
+            SELECT status FROM spot
+            WHERE spot_job_id=(?) AND task_id=(?)""",
+            (job_id, task_id)).fetchone()
+        return ManagedJobStatus(status[0]) if status else None
+
+
 def get_num_tasks(job_id: int) -> int:
     return len(_get_all_task_ids_statuses(job_id))
 
@@ -1084,8 +1096,15 @@ def get_local_log_file(job_id: int, task_id: Optional[int]) -> Optional[str]:
 
 
 def scheduler_set_waiting(job_id: int, dag_yaml_path: str, env_file_path: str,
-                          user_hash: str, priority: int) -> None:
-    """Do not call without holding the scheduler lock."""
+                          user_hash: str, priority: int) -> bool:
+    """Do not call without holding the scheduler lock.
+
+    Returns: Whether this is a recovery run or not.
+        If this is a recovery run, the job may already be in the WAITING state
+        and the update will not change the schedule_state (hence the
+        updated_count will be 0). In this case, we return True.
+        Otherwise, we return False.
+    """
     with db_utils.safe_cursor(_DB_PATH) as cursor:
         updated_count = cursor.execute(
             'UPDATE job_info SET '
@@ -1095,7 +1114,9 @@ def scheduler_set_waiting(job_id: int, dag_yaml_path: str, env_file_path: str,
             (ManagedJobScheduleState.WAITING.value, dag_yaml_path,
              env_file_path, user_hash, priority, job_id,
              ManagedJobScheduleState.INACTIVE.value)).rowcount
-        assert updated_count == 1, (job_id, updated_count)
+        # For a recovery run, the job may already be in the WAITING state.
+        assert updated_count <= 1, (job_id, updated_count)
+        return updated_count == 0
 
 
 def scheduler_set_launching(job_id: int,
diff --git a/sky/jobs/utils.py b/sky/jobs/utils.py
index 375217f0503..3cb232e27ac 100644
--- a/sky/jobs/utils.py
+++ b/sky/jobs/utils.py
@@ -176,6 +176,17 @@ def update_managed_jobs_statuses(job_id: Optional[int] = None):
     Note: we expect that job_id, if provided, refers to a nonterminal job or a
     job that has not completed its cleanup (schedule state not DONE).
     """
+    # This signal file suggests that the controller is recovering from a
+    # failure. See sky/templates/kubernetes-ray.yml.j2 for more details.
+    # When restarting the controller processes, we don't want this event to
+    # set the job status to FAILED_CONTROLLER.
+    # TODO(tian): Change this to restart the controller process. For now we
+    # disabled it when recovering because we want to avoid caveats of infinite
+    # restart of last controller process that fully occupied the controller VM.
+    if os.path.exists(
+            os.path.expanduser(
+                constants.PERSISTENT_RUN_RESTARTING_SIGNAL_FILE)):
+        return
 
     def _cleanup_job_clusters(job_id: int) -> Optional[str]:
         """Clean up clusters for a job. Returns error message if any.
diff --git a/sky/skylet/constants.py b/sky/skylet/constants.py
index 38b7a2898f3..7b7a1293e29 100644
--- a/sky/skylet/constants.py
+++ b/sky/skylet/constants.py
@@ -396,6 +396,10 @@
 # persistent through PVC. See kubernetes-ray.yml.j2.
 PERSISTENT_SETUP_SCRIPT_PATH = '~/.sky/.controller_recovery_setup_commands.sh'
 PERSISTENT_RUN_SCRIPT_DIR = '~/.sky/.controller_recovery_task_run'
+# Signal file to indicate that the controller is recovering from a failure.
+# See sky/jobs/utils.py::update_managed_jobs_statuses for more details.
+PERSISTENT_RUN_RESTARTING_SIGNAL_FILE = (
+    '~/.sky/.controller_recovery_restarting_signal')
 
 # The placeholder for the local skypilot config path in file mounts for
 # controllers.
diff --git a/sky/templates/kubernetes-ray.yml.j2 b/sky/templates/kubernetes-ray.yml.j2
index 70fa6db2c30..06049ce94a1 100644
--- a/sky/templates/kubernetes-ray.yml.j2
+++ b/sky/templates/kubernetes-ray.yml.j2
@@ -639,12 +639,14 @@ available_node_types:
                 /bin/bash --login -c "true && export OMP_NUM_THREADS=1 PYTHONWARNINGS='ignore' && {{k8s_high_availability_deployment_setup_script_path}} > /tmp/controller_recovery_setup_commands.log 2>&1"
                 echo "=== Controller setup commands completed for recovery ==="
 
+                touch {{k8s_high_availability_restarting_signal_file}}
                 for file in {{k8s_high_availability_deployment_run_script_dir}}/*; do
                   # ! Keep this aligned with `CloudVmRayBackend._execute()`
                   chmod +x $file
                   /bin/bash --login -c "true && export OMP_NUM_THREADS=1 PYTHONWARNINGS='ignore' && $file > /tmp/task_run_$(basename $file).log 2>&1"
                   echo "=== Controller task run for service (file: $file) completed for recovery ==="
                 done
+                rm {{k8s_high_availability_restarting_signal_file}}
               fi
 
               touch {{k8s_high_availability_deployment_volume_mount_path}}/k8s_container_ready
diff --git a/sky/utils/controller_utils.py b/sky/utils/controller_utils.py
index 50db197ae64..92c61686d4d 100644
--- a/sky/utils/controller_utils.py
+++ b/sky/utils/controller_utils.py
@@ -206,8 +206,7 @@ def from_type(cls, controller_type: str) -> Optional['Controllers']:
         return None
 
 
-def high_availability_specified(cluster_name: Optional[str],
-                                skip_warning: bool = True) -> bool:
+def high_availability_specified(cluster_name: Optional[str]) -> bool:
     """Check if the controller high availability is specified in user config.
     """
     controller = Controllers.from_name(cluster_name)
@@ -215,18 +214,9 @@ def high_availability_specified(cluster_name: Optional[str],
         return False
 
     if skypilot_config.loaded():
-        high_availability = skypilot_config.get_nested(
-            (controller.value.controller_type, 'controller',
-             'high_availability'), False)
-        if high_availability:
-            if controller.value.controller_type != 'serve':
-                if not skip_warning:
-                    print(f'{colorama.Fore.RED}High availability controller is'
-                          'only supported for SkyServe controller. It cannot'
-                          f'be enabled for {controller.value.name}.'
-                          f'Skipping this flag.{colorama.Style.RESET_ALL}')
-            else:
-                return True
+        return skypilot_config.get_nested((controller.value.controller_type,
+                                           'controller', 'high_availability'),
+                                          False)
     return False
 
 
diff --git a/tests/smoke_tests/smoke_tests_utils.py b/tests/smoke_tests/smoke_tests_utils.py
index 4efe72e8aa8..0d911fb18fb 100644
--- a/tests/smoke_tests/smoke_tests_utils.py
+++ b/tests/smoke_tests/smoke_tests_utils.py
@@ -783,3 +783,27 @@ def get_response_from_request_id(request_id: str) -> Any:
         return request_task.get_return_value()
     raise RuntimeError(f'Failed to get request {request_id}: '
                        f'{response.status_code} {response.text}')
+
+
+def _get_controller_pod_name(controller_name: str) -> str:
+    return (
+        'kubectl get pods -l app -o custom-columns=NAME:.metadata.name,'
+        'APP:.metadata.labels.app --no-headers | '
+        f'awk \'$2 ~ /sky-{controller_name}-controller/ {{print $1; exit}}\'')
+
+
+def kill_and_wait_controller(controller_name: str) -> str:
+    """Kill the controller pod and wait for a new one to be ready."""
+    assert controller_name in ['serve', 'jobs'
+                              ], (f'Invalid controller name: {controller_name}')
+    return (
+        f'initial_controller_pod=$({_get_controller_pod_name(controller_name)}); '
+        f'echo "Killing {controller_name} controller pod: $initial_controller_pod"; '
+        'kubectl delete pod $initial_controller_pod; '
+        f'until new_controller_pod=$({_get_controller_pod_name(controller_name)}) && '
+        '[ "$new_controller_pod" != "$initial_controller_pod" ] && '
+        'kubectl get pod $new_controller_pod | grep "1/1"; do '
+        f'  echo "Waiting for new {controller_name} controller pod..."; sleep 5; '
+        'done; '
+        f'echo "New {controller_name} controller pod ready: $new_controller_pod"'
+    )
diff --git a/tests/smoke_tests/test_managed_job.py b/tests/smoke_tests/test_managed_job.py
index 02ac94e3360..7daeadf5066 100644
--- a/tests/smoke_tests/test_managed_job.py
+++ b/tests/smoke_tests/test_managed_job.py
@@ -1163,3 +1163,46 @@ def test_managed_jobs_logs_sync_down(generic_cloud: str):
         timeout=20 * 60,
     )
     smoke_tests_utils.run_one_test(test)
+
+
+def _get_ha_kill_test(name: str, generic_cloud: str,
+                      status: sky.ManagedJobStatus) -> smoke_tests_utils.Test:
+    return smoke_tests_utils.Test(
+        f'test-managed-jobs-ha-kill-{status.value.lower()}',
+        [
+            f'sky jobs launch -n {name} --infra {generic_cloud} '
+            f'{smoke_tests_utils.LOW_RESOURCE_ARG} -y examples/managed_job.yaml -d',
+            smoke_tests_utils.
+            get_cmd_wait_until_managed_job_status_contains_matching_job_name(
+                job_name=f'{name}', job_status=[status], timeout=95),
+            smoke_tests_utils.kill_and_wait_controller('jobs'),
+            smoke_tests_utils.
+            get_cmd_wait_until_managed_job_status_contains_matching_job_name(
+                job_name=f'{name}',
+                job_status=[sky.ManagedJobStatus.SUCCEEDED],
+                timeout=335),
+            f's=$(sky jobs logs --controller -n {name} --no-follow); echo "$s"; echo "$s" | grep "Cluster launched:"',
+            rf'{smoke_tests_utils.GET_JOB_QUEUE} | grep {name} | head -n1 | grep "SUCCEEDED"',
+        ],
+        f'sky jobs cancel -y -n {name}',
+        env={
+            skypilot_config.ENV_VAR_SKYPILOT_CONFIG: 'tests/test_yamls/managed_jobs_ha_config.yaml'
+        },
+        timeout=20 * 60,
+    )
+
+
+@pytest.mark.kubernetes
+@pytest.mark.managed_jobs
+def test_managed_jobs_ha_kill_running(generic_cloud: str):
+    name = smoke_tests_utils.get_cluster_name()
+    test = _get_ha_kill_test(name, generic_cloud, sky.ManagedJobStatus.RUNNING)
+    smoke_tests_utils.run_one_test(test)
+
+
+@pytest.mark.kubernetes
+@pytest.mark.managed_jobs
+def test_managed_jobs_ha_kill_starting(generic_cloud: str):
+    name = smoke_tests_utils.get_cluster_name()
+    test = _get_ha_kill_test(name, generic_cloud, sky.ManagedJobStatus.STARTING)
+    smoke_tests_utils.run_one_test(test)
diff --git a/tests/smoke_tests/test_sky_serve.py b/tests/smoke_tests/test_sky_serve.py
index 5b526b5efc4..269e7171165 100644
--- a/tests/smoke_tests/test_sky_serve.py
+++ b/tests/smoke_tests/test_sky_serve.py
@@ -1113,7 +1113,7 @@ def test_skyserve_ha_kill_after_ready():
             f'{_SERVE_ENDPOINT_WAIT.format(name=name)}; '
             'curl $endpoint | grep "Hi, SkyPilot here"',
             # Kill controller and verify recovery
-            _kill_and_wait_controller(),
+            smoke_tests_utils.kill_and_wait_controller('serve'),
             # Verify service remains accessible after controller recovery
             _SERVE_WAIT_UNTIL_READY.format(name=name, replica_num=1),
             _check_replica_in_status(name, [(1, False, 'READY')]),
@@ -1145,7 +1145,7 @@ def test_skyserve_ha_kill_during_provision():
             f'  s=$(sky serve status {name}); '
             'done; echo "$s"',
             # Kill controller during provisioning
-            _kill_and_wait_controller(),
+            smoke_tests_utils.kill_and_wait_controller('serve'),
             # Verify service eventually becomes ready
             _SERVE_WAIT_UNTIL_READY.format(name=name, replica_num=1),
             _check_replica_in_status(name, [(1, False, 'READY')]),
@@ -1179,7 +1179,7 @@ def test_skyserve_ha_kill_during_pending():
             f'{_SERVE_STATUS_WAIT.format(name=name)}; ',
             _check_replica_in_status(name, [(1, False, 'PENDING')]),
             # Kill controller during pending
-            _kill_and_wait_controller(),
+            smoke_tests_utils.kill_and_wait_controller('serve'),
             # Verify service eventually becomes ready and accessible
             _SERVE_WAIT_UNTIL_READY.format(name=name, replica_num=1),
             _check_replica_in_status(name, [(1, False, 'READY')]),
@@ -1227,7 +1227,7 @@ def test_skyserve_ha_kill_during_shutdown():
             f'  s=$(sky serve status {name}); '
             'done; echo "$s"',
             # Kill controller during shutdown
-            _kill_and_wait_controller(),
+            smoke_tests_utils.kill_and_wait_controller('serve'),
             # Even after the pod ready, `serve status` may return `Failed to connect to serve controller, please try again later.`
             # So we need to wait for a while before checking the status again.
             'sleep 10',
@@ -1243,16 +1243,3 @@ def test_skyserve_ha_kill_during_shutdown():
         timeout=30 * 60,
         env={'SKYPILOT_CONFIG': 'tests/skyserve/high_availability/config.yaml'})
     smoke_tests_utils.run_one_test(test)
-
-
-def _kill_and_wait_controller() -> str:
-    """Kill the controller pod and wait for a new one to be ready."""
-    return (
-        'initial_controller_pod=$(kubectl get pods -l "skypilot-head-node=1" -o jsonpath="{.items[0].metadata.name}"); '
-        'echo "Killing controller pod: $initial_controller_pod"; '
-        'kubectl delete pod $initial_controller_pod; '
-        'until new_controller_pod=$(kubectl get pods -l "skypilot-head-node=1" -o jsonpath="{.items[0].metadata.name}") && '
-        '[ "$new_controller_pod" != "$initial_controller_pod" ] && kubectl get pod $new_controller_pod | grep "1/1"; do '
-        '  echo "Waiting for new controller pod..."; sleep 5; '
-        'done; '
-        'echo "New controller pod ready: $new_controller_pod"')
diff --git a/tests/test_yamls/managed_jobs_ha_config.yaml b/tests/test_yamls/managed_jobs_ha_config.yaml
new file mode 100644
index 00000000000..229cc766dc3
--- /dev/null
+++ b/tests/test_yamls/managed_jobs_ha_config.yaml
@@ -0,0 +1,6 @@
+jobs:
+  controller:
+    resources:
+      infra: kubernetes
+      cpus: 2
+    high_availability: true


Code review comments:
diff_hunk: @@ -221,8 +241,12 @@ def _run_one_task(self, task_id: int, task: 'sky.Task') -> bool:
 
             # NOTE: we do not check cluster status first because race condition
             # can occur, i.e. cluster can be down during the job status check.
-            job_status = managed_job_utils.get_job_status(
-                self._backend, cluster_name)
+            try:
+                job_status = managed_job_utils.get_job_status(
+                    self._backend, cluster_name)
+            except (exceptions.FetchClusterInfoError, exceptions.CommandError):
path: sky/jobs/controller.py
text: @user1:
CommandError should be caught by get_job_status. Is it possible to hit FetchClusterInfoError? If yes, what's the stacktrace?

@author:
Yes; if we terminate the controller on provisioning/terminating the worker cluster, and resume it, it is possible that the cluster info cannot be fetched. I encountered this today and added this except. But forget to record the stacktrace. Will do next time

@user1:
Hmm, I'm pretty sure this should be handled by get_job_status. If not it may be a bug.

diff_hunk: @@ -152,6 +152,20 @@ def _run_one_task(self, task_id: int, task: 'sky.Task') -> bool:
         Other exceptions may be raised depending on the backend.
         """
 
+        latest_task_id, status = managed_job_state.get_latest_task_id_status(
+            self._job_id)
+        is_recovery = False
path: sky/jobs/controller.py
text: @user1:
Can we call this something else? The task may still be running, in which case we don't necessarily need to "recover" - that is, recover from spot preemption.
Maybe we can call this `is_resume` or `is_restore`?

diff_hunk: @@ -171,42 +185,48 @@ def _run_one_task(self, task_id: int, task: 'sky.Task') -> bool:
             return True
         usage_lib.messages.usage.update_task_id(task_id)
         task_id_env_var = task.envs[constants.TASK_ID_ENV_VAR]
-        submitted_at = time.time()
-        if task_id == 0:
-            submitted_at = backend_utils.get_timestamp_from_run_timestamp(
-                self._backend.run_timestamp)
         assert task.name is not None, task
         cluster_name = managed_job_utils.generate_managed_job_cluster_name(
             task.name, self._job_id)
         self._strategy_executor = recovery_strategy.StrategyExecutor.make(
             cluster_name, self._backend, task, self._job_id, task_id)
-        managed_job_state.set_starting(
-            self._job_id,
-            task_id,
-            self._backend.run_timestamp,
-            submitted_at,
-            resources_str=backend_utils.get_task_resources_str(
-                task, is_managed_job=True),
-            specs={
-                'max_restarts_on_errors':
-                    self._strategy_executor.max_restarts_on_errors
-            },
-            callback_func=callback_func)
-        logger.info(
-            f'Submitted managed job {self._job_id} (task: {task_id}, name: '
-            f'{task.name!r}); {constants.TASK_ID_ENV_VAR}: {task_id_env_var}')
-
-        logger.info('Started monitoring.')
-
-        remote_job_submitted_at = self._strategy_executor.launch()
-        assert remote_job_submitted_at is not None, remote_job_submitted_at
+        if not is_recovery:
+            submitted_at = time.time()
+            if task_id == 0:
+                submitted_at = backend_utils.get_timestamp_from_run_timestamp(
+                    self._backend.run_timestamp)
+            managed_job_state.set_starting(
+                self._job_id,
+                task_id,
+                self._backend.run_timestamp,
+                submitted_at,
+                resources_str=backend_utils.get_task_resources_str(
+                    task, is_managed_job=True),
+                specs={
+                    'max_restarts_on_errors':
+                        self._strategy_executor.max_restarts_on_errors
+                },
+                callback_func=callback_func)
+            logger.info(f'Submitted managed job {self._job_id} '
+                        f'(task: {task_id}, name: {task.name!r}); '
+                        f'{constants.TASK_ID_ENV_VAR}: {task_id_env_var}')
+
+            logger.info('Started monitoring.')
+
+            remote_job_submitted_at = self._strategy_executor.launch()
+            assert remote_job_submitted_at is not None, remote_job_submitted_at
 
-        managed_job_state.set_started(job_id=self._job_id,
-                                      task_id=task_id,
-                                      start_time=remote_job_submitted_at,
-                                      callback_func=callback_func)
+            managed_job_state.set_started(job_id=self._job_id,
+                                          task_id=task_id,
+                                          start_time=remote_job_submitted_at,
+                                          callback_func=callback_func)
 
         while True:
+            if is_recovery:
+                last_status = managed_job_state.get_job_status_with_task_id(
+                    job_id=self._job_id, task_id=task_id)
+                if last_status is not None and last_status.is_terminal():
+                    return True
path: sky/jobs/controller.py
text: @user1:
Should return True only if this task is SUCCEEDED, so we don't proceed to later tasks.
```suggestion
                    return last_status == managed_job_state.SUCCEEDED
```

diff_hunk: @@ -639,12 +639,14 @@ available_node_types:
                 /bin/bash --login -c "true && export OMP_NUM_THREADS=1 PYTHONWARNINGS='ignore' && {{k8s_high_availability_deployment_setup_script_path}} > /tmp/controller_recovery_setup_commands.log 2>&1"
                 echo "=== Controller setup commands completed for recovery ==="
 
+                touch {{k8s_high_availability_restarting_signal_file}}
                 for file in {{k8s_high_availability_deployment_run_script_dir}}/*; do
                   # ! Keep this aligned with `CloudVmRayBackend._execute()`
                   chmod +x $file
                   /bin/bash --login -c "true && export OMP_NUM_THREADS=1 PYTHONWARNINGS='ignore' && $file > /tmp/task_run_$(basename $file).log 2>&1"
path: sky/templates/kubernetes-ray.yml.j2
text: @user1:
This logic may run a lot of things if the jobs controller previously had many jobs. We should make sure it will scale well.

@author:
Good point! Added a todo here.