skypilot-nightly 1.0.0.dev20250509__py3-none-any.whl → 1.0.0.dev20251107__py3-none-any.whl

sky/jobs/recovery_strategy.py

@@ -5,23 +5,31 @@ In the YAML file, the user can specify the strategy to use for managed jobs.
 resources:
     job_recovery: EAGER_NEXT_REGION
 """
-import time
+import asyncio
+import logging
+import os
 import traceback
 import typing
-from typing import Optional
+from typing import Optional, Set
 
-import sky
 from sky import backends
+from sky import dag as dag_lib
 from sky import exceptions
-from sky import execution
 from sky import global_user_state
 from sky import sky_logging
+from sky import skypilot_config
 from sky.backends import backend_utils
+from sky.client import sdk
 from sky.jobs import scheduler
+from sky.jobs import state
 from sky.jobs import utils as managed_job_utils
+from sky.serve import serve_utils
+from sky.skylet import constants
 from sky.skylet import job_lib
 from sky.usage import usage_lib
 from sky.utils import common_utils
+from sky.utils import context_utils
+from sky.utils import env_options
 from sky.utils import registry
 from sky.utils import status_lib
 from sky.utils import ux_utils
@@ -39,7 +47,14 @@ MAX_JOB_CHECKING_RETRY = 10
 # Minutes to job cluster autodown. This should be significantly larger than
 # managed_job_utils.JOB_STATUS_CHECK_GAP_SECONDS, to avoid tearing down the
 # cluster before its status can be updated by the job controller.
-_AUTODOWN_MINUTES = 5
+_AUTODOWN_MINUTES = 10
+
+ENV_VARS_TO_CLEAR = [
+    skypilot_config.ENV_VAR_SKYPILOT_CONFIG,
+    constants.USER_ID_ENV_VAR,
+    constants.USER_ENV_VAR,
+    env_options.Options.SHOW_DEBUG_INFO.env_key,
+]
 
 
 class StrategyExecutor:
@@ -47,29 +62,65 @@ class StrategyExecutor:
 
     RETRY_INIT_GAP_SECONDS = 60
 
-    def __init__(self, cluster_name: str, backend: 'backends.Backend',
-                 task: 'task_lib.Task', max_restarts_on_errors: int,
-                 job_id: int) -> None:
+    def __init__(
+        self,
+        cluster_name: Optional[str],
+        backend: 'backends.Backend',
+        task: 'task_lib.Task',
+        max_restarts_on_errors: int,
+        job_id: int,
+        task_id: int,
+        pool: Optional[str],
+        starting: Set[int],
+        starting_lock: asyncio.Lock,
+        starting_signal: asyncio.Condition,
+    ) -> None:
         """Initialize the strategy executor.
 
         Args:
             cluster_name: The name of the cluster.
             backend: The backend to use. Only CloudVMRayBackend is supported.
             task: The task to execute.
+            max_restarts_on_errors: Maximum number of restarts on errors.
+            job_id: The ID of the job.
+            task_id: The ID of the task.
+            starting: Set of job IDs that are currently starting.
+            starting_lock: Lock to synchronize starting jobs.
+            starting_signal: Condition to signal when a job can start.
         """
         assert isinstance(backend, backends.CloudVmRayBackend), (
             'Only CloudVMRayBackend is supported.')
-        self.dag = sky.Dag()
+        self.dag = dag_lib.Dag()
         self.dag.add(task)
+        # For jobs submitted to a pool, the cluster name might change after each
+        # recovery. Initially this is set to an empty string to indicate that no
+        # cluster is assigned yet, and in `_launch`, it will be set to one of
+        # the cluster names in the pool.
         self.cluster_name = cluster_name
         self.backend = backend
         self.max_restarts_on_errors = max_restarts_on_errors
         self.job_id = job_id
+        self.task_id = task_id
+        self.pool = pool
         self.restart_cnt_on_failure = 0
+        self.job_id_on_pool_cluster: Optional[int] = None
+        self.starting = starting
+        self.starting_lock = starting_lock
+        self.starting_signal = starting_signal
 
     @classmethod
-    def make(cls, cluster_name: str, backend: 'backends.Backend',
-             task: 'task_lib.Task', job_id: int) -> 'StrategyExecutor':
+    def make(
+        cls,
+        cluster_name: Optional[str],
+        backend: 'backends.Backend',
+        task: 'task_lib.Task',
+        job_id: int,
+        task_id: int,
+        pool: Optional[str],
+        starting: Set[int],
+        starting_lock: asyncio.Lock,
+        starting_signal: asyncio.Condition,
+    ) -> 'StrategyExecutor':
         """Create a strategy from a task."""
 
         resource_list = list(task.resources)
@@ -86,8 +137,11 @@ class StrategyExecutor:
         # original task.resources
         task.set_resources(type(task.resources)(new_resources_list))
         if isinstance(job_recovery, dict):
-            job_recovery_name = job_recovery.pop(
+            name = job_recovery.pop(
                 'strategy', registry.JOBS_RECOVERY_STRATEGY_REGISTRY.default)
+            assert name is None or isinstance(name, str), (
+                name, 'The job recovery strategy name must be a string or None')
+            job_recovery_name: Optional[str] = name
             max_restarts_on_errors = job_recovery.pop('max_restarts_on_errors',
                                                       0)
         else:
@@ -97,9 +151,11 @@ class StrategyExecutor:
                                  from_str(job_recovery_name))
         assert job_recovery_strategy is not None, job_recovery_name
         return job_recovery_strategy(cluster_name, backend, task,
-                                     max_restarts_on_errors, job_id)
+                                     max_restarts_on_errors, job_id, task_id,
+                                     pool, starting, starting_lock,
+                                     starting_signal)
 
-    def launch(self) -> float:
+    async def launch(self) -> float:
         """Launch the cluster for the first time.
 
         It can fail if resource is not available. Need to check the cluster
@@ -111,11 +167,11 @@ class StrategyExecutor:
         Raises: Please refer to the docstring of self._launch().
         """
 
-        job_submit_at = self._launch(max_retry=None)
+        job_submit_at = await self._launch(max_retry=None)
         assert job_submit_at is not None
         return job_submit_at
 
-    def recover(self) -> float:
+    async def recover(self) -> float:
         """Relaunch the cluster after failure and wait until job starts.
 
         When recover() is called the cluster should be in STOPPED status (i.e.
@@ -125,12 +181,12 @@ class StrategyExecutor:
         """
         raise NotImplementedError
 
-    def _try_cancel_all_jobs(self):
-        from sky import core  # pylint: disable=import-outside-toplevel
-
-        handle = global_user_state.get_handle_from_cluster_name(
-            self.cluster_name)
-        if handle is None:
+    async def _try_cancel_jobs(self):
+        if self.cluster_name is None:
+            return
+        handle = await context_utils.to_thread(
+            global_user_state.get_handle_from_cluster_name, self.cluster_name)
+        if handle is None or self.pool is not None:
             return
         try:
             usage_lib.messages.usage.set_internal()
@@ -153,9 +209,22 @@ class StrategyExecutor:
             # should be functional with the `_try_cancel_if_cluster_is_init`
             # flag, i.e. it sends the cancel signal to the head node, which will
             # then kill the user process on remaining worker nodes.
-            core.cancel(cluster_name=self.cluster_name,
-                        all=True,
-                        _try_cancel_if_cluster_is_init=True)
+            # Only cancel the corresponding job for worker pool.
+            if self.pool is None:
+                kwargs = dict(all=True)
+            else:
+                kwargs = dict(job_ids=[self.job_id_on_pool_cluster])
+            request_id = await context_utils.to_thread(
+                sdk.cancel,
+                cluster_name=self.cluster_name,
+                **kwargs,
+                _try_cancel_if_cluster_is_init=True,
+            )
+            logger.debug(f'sdk.cancel request ID: {request_id}')
+            await context_utils.to_thread(
+                sdk.get,
+                request_id,
+            )
         except Exception as e:  # pylint: disable=broad-except
             logger.info('Failed to cancel the job on the cluster. The cluster '
                         'might be already down or the head node is preempted.'
@@ -163,25 +232,26 @@ class StrategyExecutor:
                         f'{common_utils.format_exception(e)}\n'
                         'Terminating the cluster explicitly to ensure no '
                         'remaining job process interferes with recovery.')
-            managed_job_utils.terminate_cluster(self.cluster_name)
+            await context_utils.to_thread(self._cleanup_cluster)
 
-    def _wait_until_job_starts_on_cluster(self) -> Optional[float]:
+    async def _wait_until_job_starts_on_cluster(self) -> Optional[float]:
         """Wait for MAX_JOB_CHECKING_RETRY times until job starts on the cluster
 
         Returns:
             The timestamp of when the job is submitted, or None if failed to
             submit.
         """
+        assert self.cluster_name is not None
         status = None
         job_checking_retry_cnt = 0
         while job_checking_retry_cnt < MAX_JOB_CHECKING_RETRY:
             # Avoid the infinite loop, if any bug happens.
             job_checking_retry_cnt += 1
             try:
-                cluster_status, _ = (
-                    backend_utils.refresh_cluster_status_handle(
-                        self.cluster_name,
-                        force_refresh_statuses=set(status_lib.ClusterStatus)))
+                cluster_status, _ = (await context_utils.to_thread(
+                    backend_utils.refresh_cluster_status_handle,
+                    self.cluster_name,
+                    force_refresh_statuses=set(status_lib.ClusterStatus)))
             except Exception as e:  # pylint: disable=broad-except
                 # If any unexpected error happens, retry the job checking
                 # loop.
@@ -201,8 +271,10 @@ class StrategyExecutor:
                 break
 
             try:
-                status = managed_job_utils.get_job_status(
-                    self.backend, self.cluster_name)
+                status = await managed_job_utils.get_job_status(
+                    self.backend,
+                    self.cluster_name,
+                    job_id=self.job_id_on_pool_cluster)
             except Exception as e:  # pylint: disable=broad-except
                 # If any unexpected error happens, retry the job checking
                 # loop.
@@ -217,8 +289,12 @@ class StrategyExecutor:
             # Check the job status until it is not in initialized status
             if status is not None and status > job_lib.JobStatus.INIT:
                 try:
-                    job_submitted_at = managed_job_utils.get_job_timestamp(
-                        self.backend, self.cluster_name, get_end_time=False)
+                    job_submitted_at = await context_utils.to_thread(
+                        managed_job_utils.get_job_timestamp,
+                        self.backend,
+                        self.cluster_name,
+                        self.job_id_on_pool_cluster,
+                        get_end_time=False)
                     return job_submitted_at
                 except Exception as e:  # pylint: disable=broad-except
                     # If we failed to get the job timestamp, we will retry
@@ -227,12 +303,20 @@ class StrategyExecutor:
                                 'the job start timestamp. Retrying.')
                     continue
             # Wait for the job to be started
-            time.sleep(managed_job_utils.JOB_STARTED_STATUS_CHECK_GAP_SECONDS)
+            await asyncio.sleep(
+                managed_job_utils.JOB_STARTED_STATUS_CHECK_GAP_SECONDS)
         return None
 
-    def _launch(self,
-                max_retry: Optional[int] = 3,
-                raise_on_failure: bool = True) -> Optional[float]:
+    def _cleanup_cluster(self) -> None:
+        if self.cluster_name is None:
+            return
+        if self.pool is None:
+            managed_job_utils.terminate_cluster(self.cluster_name)
+
+    async def _launch(self,
+                      max_retry: Optional[int] = 3,
+                      raise_on_failure: bool = True,
+                      recovery: bool = False) -> Optional[float]:
         """Implementation of launch().
 
         The function will wait until the job starts running, but will leave the
@@ -272,98 +356,234 @@ class StrategyExecutor:
         backoff = common_utils.Backoff(self.RETRY_INIT_GAP_SECONDS)
         while True:
             retry_cnt += 1
-            with scheduler.scheduled_launch(self.job_id):
-                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).
-                    execution.launch(
-                        self.dag,
-                        cluster_name=self.cluster_name,
-                        # We expect to tear down the cluster as soon as the job
-                        # is finished. However, in case the controller dies, set
-                        # autodown to try and avoid a resource leak.
-                        idle_minutes_to_autostop=_AUTODOWN_MINUTES,
-                        down=True,
-                        _is_launched_by_jobs_controller=True)
-                    logger.info('Managed job cluster launched.')
-                except (exceptions.InvalidClusterNameError,
-                        exceptions.NoCloudAccessError,
-                        exceptions.ResourcesMismatchError) as e:
-                    logger.error('Failure happened before provisioning. '
-                                 f'{common_utils.format_exception(e)}')
-                    if raise_on_failure:
-                        raise exceptions.ProvisionPrechecksError(reasons=[e])
-                    return None
-                except exceptions.ResourcesUnavailableError as e:
-                    # This is raised when the launch fails due to prechecks or
-                    # after failing over through all the candidates.
-                    # Please refer to the docstring of `sky.launch` for more
-                    # details of how the exception will be structured.
-                    if not any(
-                            isinstance(err,
-                                       exceptions.ResourcesUnavailableError)
-                            for err in e.failover_history):
-                        # _launch() (this function) should fail/exit directly,
-                        # if none of the failover reasons were because of
-                        # resource unavailability or no failover was attempted
-                        # (the optimizer cannot find feasible resources for
-                        # requested resources), i.e., e.failover_history is
-                        # empty. Failing directly avoids the infinite loop of
-                        # retrying the launch when, e.g., an invalid cluster
-                        # name is used and --retry-until-up is specified.
-                        reasons = (e.failover_history
-                                   if e.failover_history else [e])
-                        reasons_str = '; '.join(
-                            common_utils.format_exception(err)
-                            for err in reasons)
-                        logger.error(
-                            'Failure happened before provisioning. Failover '
-                            f'reasons: {reasons_str}')
+            try:
+                async with scheduler.scheduled_launch(
+                        self.job_id,
+                        self.starting,
+                        self.starting_lock,
+                        self.starting_signal,
+                ):
+                    # The job state may have been PENDING during backoff -
+                    # update to STARTING or RECOVERING.
+                    # On the first attempt (when retry_cnt is 1), we should
+                    # already be in STARTING or RECOVERING.
+                    if retry_cnt > 1:
+                        await state.set_restarting_async(
+                            self.job_id, self.task_id, recovery)
+                    try:
+                        usage_lib.messages.usage.set_internal()
+                        if self.pool is None:
+                            assert self.cluster_name is not None
+
+                            # sdk.launch will implicitly start the API server,
+                            # but then the API server will inherit the current
+                            # env vars/user, which we may not want.
+                            # Instead, clear env vars here and call api_start
+                            # explicitly.
+                            vars_to_restore = {}
+                            try:
+                                for env_var in ENV_VARS_TO_CLEAR:
+                                    vars_to_restore[env_var] = os.environ.pop(
+                                        env_var, None)
+                                    logger.debug('Cleared env var: '
+                                                 f'{env_var}')
+                                logger.debug('Env vars for api_start: '
+                                             f'{os.environ}')
+                                await context_utils.to_thread(sdk.api_start)
+                                logger.info('API server started.')
+                            finally:
+                                for env_var, value in vars_to_restore.items():
+                                    if value is not None:
+                                        logger.debug('Restored env var: '
+                                                     f'{env_var}: {value}')
+                                        os.environ[env_var] = value
+
+                            request_id = None
+                            try:
+                                request_id = await context_utils.to_thread(
+                                    sdk.launch,
+                                    self.dag,
+                                    cluster_name=self.cluster_name,
+                                    # We expect to tear down the cluster as soon
+                                    # as the job is finished. However, in case
+                                    # the controller dies, we may end up with a
+                                    # resource leak.
+                                    # Ideally, we should autodown to be safe,
+                                    # but it's fine to disable it for now, as
+                                    # Nebius doesn't support autodown yet.
+                                    # TODO(kevin): set down=True once Nebius
+                                    # supports autodown.
+                                    # idle_minutes_to_autostop=(
+                                    #     _AUTODOWN_MINUTES),
+                                    # down=True,
+                                    _is_launched_by_jobs_controller=True,
+                                )
+                                logger.debug('sdk.launch request ID: '
+                                             f'{request_id}')
+                                await context_utils.to_thread(
+                                    sdk.stream_and_get,
+                                    request_id,
+                                )
+                            except asyncio.CancelledError:
+                                if request_id:
+                                    req = await context_utils.to_thread(
+                                        sdk.api_cancel, request_id)
+                                    logger.debug('sdk.api_cancel request '
+                                                 f'ID: {req}')
+                                    try:
+                                        await context_utils.to_thread(
+                                            sdk.get, req)
+                                    except Exception as e:  # pylint: disable=broad-except
+                                        # we must still return a CancelledError
+                                        logger.error(
+                                            f'Failed to cancel the job: {e}')
+                                raise
+                            logger.info('Managed job cluster launched.')
+                        else:
+                            self.cluster_name = await (context_utils.to_thread(
+                                serve_utils.get_next_cluster_name, self.pool,
+                                self.job_id))
+                            if self.cluster_name is None:
+                                raise exceptions.NoClusterLaunchedError(
+                                    'No cluster name found in the pool.')
+                            request_id = None
+                            try:
+                                request_id = await context_utils.to_thread(
+                                    sdk.exec,
+                                    self.dag,
+                                    cluster_name=self.cluster_name,
+                                )
+                                logger.debug('sdk.exec request ID: '
+                                             f'{request_id}')
+                                job_id_on_pool_cluster, _ = (
+                                    await context_utils.to_thread(
+                                        sdk.get, request_id))
+                            except asyncio.CancelledError:
+                                if request_id:
+                                    req = await context_utils.to_thread(
+                                        sdk.api_cancel, request_id)
+                                    logger.debug('sdk.api_cancel request '
+                                                 f'ID: {req}')
+                                    try:
+                                        await context_utils.to_thread(
+                                            sdk.get, req)
+                                    except Exception as e:  # pylint: disable=broad-except
+                                        # we must still return a CancelledError
+                                        logger.error(
+                                            f'Failed to cancel the job: {e}')
+                                raise
+                            assert job_id_on_pool_cluster is not None, (
+                                self.cluster_name, self.job_id)
+                            self.job_id_on_pool_cluster = job_id_on_pool_cluster
+                            await state.set_job_id_on_pool_cluster_async(
+                                self.job_id, job_id_on_pool_cluster)
+                        logger.info('Managed job cluster launched.')
+                    except (exceptions.InvalidClusterNameError,
+                            exceptions.NoCloudAccessError,
+                            exceptions.ResourcesMismatchError,
+                            exceptions.StorageSpecError,
+                            exceptions.StorageError) as e:
+                        logger.error('Failure happened before provisioning. '
+                                     f'{common_utils.format_exception(e)}')
                         if raise_on_failure:
-                            raise exceptions.ProvisionPrechecksError(reasons)
-                        return None
-                    logger.info('Failed to launch a cluster with error: '
-                                f'{common_utils.format_exception(e)})')
-                except Exception as e:  # pylint: disable=broad-except
-                    # If the launch fails, it will be recovered by the following
-                    # code.
-                    logger.info('Failed to launch a cluster with error: '
-                                f'{common_utils.format_exception(e)})')
-                    with ux_utils.enable_traceback():
-                        logger.info(f'  Traceback: {traceback.format_exc()}')
-                else:  # No exception, the launch succeeds.
-                    # At this point, a sky.launch() has succeeded. Cluster may
-                    # be UP (no preemption since) or DOWN (newly preempted).
-                    job_submitted_at = self._wait_until_job_starts_on_cluster()
-                    if job_submitted_at is not None:
-                        return job_submitted_at
-                    # The job fails to start on the cluster, retry the launch.
-                    # TODO(zhwu): log the unexpected error to usage collection
-                    # for future debugging.
-                    logger.info(
-                        'Failed to successfully submit the job to the '
-                        'launched cluster, due to unexpected submission errors '
-                        'or the cluster being preempted during job submission.')
-
-                # If we get here, the launch did not succeed. Tear down the
-                # cluster and retry.
-                managed_job_utils.terminate_cluster(self.cluster_name)
-                if max_retry is not None and retry_cnt >= max_retry:
-                    # Retry forever if max_retry is None.
-                    if raise_on_failure:
-                        with ux_utils.print_exception_no_traceback():
-                            raise exceptions.ManagedJobReachedMaxRetriesError(
-                                'Resources unavailable: failed to launch '
-                                f'clusters after {max_retry} retries.')
-                    else:
+                            raise exceptions.ProvisionPrechecksError(
+                                reasons=[e])
                         return None
-            # Exit the scheduled_launch context so that the scheulde state is
-            # ALIVE during the backoff. This allows other jobs to launch.
-            gap_seconds = backoff.current_backoff()
-            logger.info('Retrying to launch the cluster in '
-                        f'{gap_seconds:.1f} seconds.')
-            time.sleep(gap_seconds)
+                    except exceptions.ResourcesUnavailableError as e:
+                        # This is raised when the launch fails due to prechecks
+                        # or after failing over through all the candidates.
+                        # Please refer to the docstring of `sky.launch` for more
+                        # details of how the exception will be structured.
+                        if not any(
+                                isinstance(err,
+                                           exceptions.ResourcesUnavailableError)
+                                for err in e.failover_history):
+                            # _launch() (this function) should fail/exit
+                            # directly, if none of the failover reasons were
+                            # because of resource unavailability or no failover
+                            # was attempted (the optimizer cannot find feasible
+                            # resources for requested resources), i.e.,
+                            # e.failover_history is empty. Failing directly
+                            # avoids the infinite loop of retrying the launch
+                            # when, e.g., an invalid cluster name is used and
+                            # --retry-until-up is specified.
+                            reasons = (e.failover_history
+                                       if e.failover_history else [e])
+                            reasons_str = '; '.join(
+                                common_utils.format_exception(err)
+                                for err in reasons)
+                            logger.error(
+                                'Failure happened before provisioning. '
+                                f'Failover reasons: {reasons_str}')
+                            if raise_on_failure:
+                                raise exceptions.ProvisionPrechecksError(
+                                    reasons)
+                            return None
+                        logger.info('Failed to launch a cluster with error: '
+                                    f'{common_utils.format_exception(e)})')
+                    except Exception as e:  # pylint: disable=broad-except
+                        # If the launch fails, it will be recovered by the
+                        # following code.
+                        logger.info('Failed to launch a cluster with error: '
+                                    f'{common_utils.format_exception(e)})')
+                        with ux_utils.enable_traceback():
+                            logger.info(
+                                f'  Traceback: {traceback.format_exc()}')
+                    else:  # No exception, the launch succeeds.
+                        # At this point, a sky.launch() has succeeded. Cluster
+                        # may be UP (no preemption since) or DOWN (newly
+                        # preempted).
+                        job_submitted_at = await (
+                            self._wait_until_job_starts_on_cluster())
+                        if job_submitted_at is not None:
+                            return job_submitted_at
+                        # The job fails to start on the cluster, retry the
+                        # launch.
+                        # TODO(zhwu): log the unexpected error to usage
+                        # collection for future debugging.
+                        logger.info(
+                            'Failed to successfully submit the job to the '
+                            'launched cluster, due to unexpected submission '
+                            'errors or the cluster being preempted during '
+                            'job submission.')
+
+                    # If we get here, the launch did not succeed. Tear down the
+                    # cluster and retry.
+                    await context_utils.to_thread(self._cleanup_cluster)
+                    if max_retry is not None and retry_cnt >= max_retry:
+                        # Retry forever if max_retry is None.
+                        if raise_on_failure:
+                            with ux_utils.print_exception_no_traceback():
+                                raise (
+                                    exceptions.ManagedJobReachedMaxRetriesError(
+                                        'Resources unavailable: failed to '
+                                        f'launch clusters after {max_retry} '
+                                        'retries.'))
+                        else:
+                            return None
+
+                    # Raise NoClusterLaunchedError to indicate that the job is
+                    # in retry backoff. This will trigger special handling in
+                    # scheduler.schedule_launched().
+                    # We will exit the scheduled_launch context so that the
+                    # schedule state is ALIVE_BACKOFF during the backoff. This
+                    # allows other jobs to launch.
+                    raise exceptions.NoClusterLaunchedError()
+
+            except exceptions.NoClusterLaunchedError:
+                # Update the status to PENDING during backoff.
+                await state.set_backoff_pending_async(self.job_id, self.task_id)
+                # Calculate the backoff time and sleep.
+                gap_seconds = (backoff.current_backoff()
+                               if self.pool is None else 1)
+                logger.info('Retrying to launch the cluster in '
+                            f'{gap_seconds:.1f} seconds.')
+                await asyncio.sleep(gap_seconds)
+                continue
+            else:
+                # The inner loop should either return or throw
+                # NoClusterLaunchedError.
+                assert False, 'Unreachable'
 
     def should_restart_on_failure(self) -> bool:
         """Increments counter & checks if job should be restarted on a failure.
@@ -384,24 +604,38 @@ class FailoverStrategyExecutor(StrategyExecutor):
 
     _MAX_RETRY_CNT = 240  # Retry for 4 hours.
 
-    def __init__(self, cluster_name: str, backend: 'backends.Backend',
-                 task: 'task_lib.Task', max_restarts_on_errors: int,
-                 job_id: int) -> None:
+    def __init__(
+        self,
+        cluster_name: Optional[str],
+        backend: 'backends.Backend',
+        task: 'task_lib.Task',
+        max_restarts_on_errors: int,
+        job_id: int,
+        task_id: int,
+        pool: Optional[str],
+        starting: Set[int],
+        starting_lock: asyncio.Lock,
+        starting_signal: asyncio.Condition,
+    ) -> None:
         super().__init__(cluster_name, backend, task, max_restarts_on_errors,
-                         job_id)
+                         job_id, task_id, pool, starting, starting_lock,
+                         starting_signal)
         # Note down the cloud/region of the launched cluster, so that we can
         # first retry in the same cloud/region. (Inside recover() we may not
         # rely on cluster handle, as it can be None if the cluster is
         # preempted.)
         self._launched_resources: Optional['resources.Resources'] = None
 
-    def _launch(self,
-                max_retry: Optional[int] = 3,
-                raise_on_failure: bool = True) -> Optional[float]:
-        job_submitted_at = super()._launch(max_retry, raise_on_failure)
-        if job_submitted_at is not None:
+    async def _launch(self,
+                      max_retry: Optional[int] = 3,
+                      raise_on_failure: bool = True,
+                      recovery: bool = False) -> Optional[float]:
+        job_submitted_at = await super()._launch(max_retry, raise_on_failure,
+                                                 recovery)
+        if job_submitted_at is not None and self.cluster_name is not None:
             # Only record the cloud/region if the launch is successful.
-            handle = global_user_state.get_handle_from_cluster_name(
+            handle = await context_utils.to_thread(
+                global_user_state.get_handle_from_cluster_name,
                 self.cluster_name)
             assert isinstance(handle, backends.CloudVmRayResourceHandle), (
                 'Cluster should be launched.', handle)
@@ -411,7 +645,7 @@ class FailoverStrategyExecutor(StrategyExecutor):
             self._launched_resources = None
         return job_submitted_at
 
-    def recover(self) -> float:
+    async def recover(self) -> float:
         # 1. Cancel the jobs and launch the cluster with the STOPPED status,
         #    so that it will try on the current region first until timeout.
         # 2. Tear down the cluster, if the step 1 failed to launch the cluster.
@@ -419,7 +653,7 @@ class FailoverStrategyExecutor(StrategyExecutor):
         #    original user specification.
 
         # Step 1
-        self._try_cancel_all_jobs()
+        await self._try_cancel_jobs()
 
         while True:
             # Add region constraint to the task, to retry on the same region
@@ -433,7 +667,8 @@ class FailoverStrategyExecutor(StrategyExecutor):
                     cloud=launched_cloud, region=launched_region, zone=None)
                 task.set_resources({new_resources})
                 # Not using self.launch to avoid the retry until up logic.
-                job_submitted_at = self._launch(raise_on_failure=False)
+                job_submitted_at = await self._launch(raise_on_failure=False,
+                                                      recovery=True)
                 # Restore the original dag, i.e. reset the region constraint.
                 task.set_resources(original_resources)
                 if job_submitted_at is not None:
@@ -442,20 +677,21 @@ class FailoverStrategyExecutor(StrategyExecutor):
             # Step 2
             logger.debug('Terminating unhealthy cluster and reset cloud '
                          'region.')
-            managed_job_utils.terminate_cluster(self.cluster_name)
+            await context_utils.to_thread(self._cleanup_cluster)
 
             # Step 3
             logger.debug('Relaunch the cluster  without constraining to prior '
                          'cloud/region.')
             # Not using self.launch to avoid the retry until up logic.
-            job_submitted_at = self._launch(max_retry=self._MAX_RETRY_CNT,
-                                            raise_on_failure=False)
+            job_submitted_at = await self._launch(max_retry=self._MAX_RETRY_CNT,
+                                                  raise_on_failure=False,
+                                                  recovery=True)
             if job_submitted_at is None:
                 # Failed to launch the cluster.
                 gap_seconds = self.RETRY_INIT_GAP_SECONDS
                 logger.info('Retrying to recover the cluster in '
                             f'{gap_seconds:.1f} seconds.')
-                time.sleep(gap_seconds)
+                await asyncio.sleep(gap_seconds)
                 continue
 
             return job_submitted_at
@@ -487,7 +723,7 @@ class EagerFailoverStrategyExecutor(FailoverStrategyExecutor):
                                                   -> R1Z1 (success)
     """
 
-    def recover(self) -> float:
+    async def recover(self) -> float:
         # 1. Terminate the current cluster
         # 2. Launch again by explicitly blocking the previously launched region
         # (this will failover through the entire search space except the
@@ -500,7 +736,7 @@ class EagerFailoverStrategyExecutor(FailoverStrategyExecutor):
 
         # Step 1
         logger.debug('Terminating unhealthy cluster and reset cloud region.')
-        managed_job_utils.terminate_cluster(self.cluster_name)
+        await context_utils.to_thread(self._cleanup_cluster)
 
         # Step 2
         logger.debug('Relaunch the cluster skipping the previously launched '
@@ -521,7 +757,8 @@ class EagerFailoverStrategyExecutor(FailoverStrategyExecutor):
                                              region=launched_region)
                 }
                 # Not using self.launch to avoid the retry until up logic.
-                job_submitted_at = self._launch(raise_on_failure=False)
+                job_submitted_at = await self._launch(raise_on_failure=False,
+                                                      recovery=True)
                 task.blocked_resources = None
                 if job_submitted_at is not None:
                     return job_submitted_at
@@ -531,14 +768,23 @@ class EagerFailoverStrategyExecutor(FailoverStrategyExecutor):
             logger.debug('Relaunch the cluster without constraining to prior '
                          'cloud/region.')
             # Not using self.launch to avoid the retry until up logic.
-            job_submitted_at = self._launch(max_retry=self._MAX_RETRY_CNT,
-                                            raise_on_failure=False)
+            job_submitted_at = await self._launch(max_retry=self._MAX_RETRY_CNT,
+                                                  raise_on_failure=False,
+                                                  recovery=True)
             if job_submitted_at is None:
                 # Failed to launch the cluster.
                 gap_seconds = self.RETRY_INIT_GAP_SECONDS
                 logger.info('Retrying to recover the cluster in '
                             f'{gap_seconds:.1f} seconds.')
-                time.sleep(gap_seconds)
+                await asyncio.sleep(gap_seconds)
                 continue
 
             return job_submitted_at
+
+
+def _get_logger_file(file_logger: logging.Logger) -> Optional[str]:
+    """Gets the file path that the logger writes to."""
+    for handler in file_logger.handlers:
+        if isinstance(handler, logging.FileHandler):
+            return handler.baseFilename
+    return None