skypilot-nightly 1.0.0.dev2024053101__py3-none-any.whl → 1.0.0.dev2025022801__py3-none-any.whl

sky/jobs/recovery_strategy.py

@@ -13,50 +13,33 @@ from typing import Optional
 import sky
 from sky import backends
 from sky import exceptions
+from sky import execution
 from sky import global_user_state
 from sky import sky_logging
-from sky import status_lib
 from sky.backends import backend_utils
+from sky.jobs import scheduler
 from sky.jobs import utils as managed_job_utils
 from sky.skylet import job_lib
 from sky.usage import usage_lib
 from sky.utils import common_utils
+from sky.utils import registry
+from sky.utils import status_lib
 from sky.utils import ux_utils
 
 if typing.TYPE_CHECKING:
+    from sky import resources
     from sky import task as task_lib
 
 logger = sky_logging.init_logger(__name__)
 
-RECOVERY_STRATEGIES = {}
-DEFAULT_RECOVERY_STRATEGY = None
-
 # Waiting time for job from INIT/PENDING to RUNNING
 # 10 * JOB_STARTED_STATUS_CHECK_GAP_SECONDS = 10 * 5 = 50 seconds
 MAX_JOB_CHECKING_RETRY = 10
 
-
-def terminate_cluster(cluster_name: str, max_retry: int = 3) -> None:
-    """Terminate the cluster."""
-    retry_cnt = 0
-    while True:
-        try:
-            usage_lib.messages.usage.set_internal()
-            sky.down(cluster_name)
-            return
-        except ValueError:
-            # The cluster is already down.
-            return
-        except Exception as e:  # pylint: disable=broad-except
-            retry_cnt += 1
-            if retry_cnt >= max_retry:
-                raise RuntimeError(
-                    f'Failed to terminate the cluster {cluster_name}.') from e
-            logger.error(
-                f'Failed to terminate the cluster {cluster_name}. Retrying.'
-                f'Details: {common_utils.format_exception(e)}')
-            with ux_utils.enable_traceback():
-                logger.error(f'  Traceback: {traceback.format_exc()}')
+# 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
 
 
 class StrategyExecutor:
@@ -65,14 +48,14 @@ class StrategyExecutor:
     RETRY_INIT_GAP_SECONDS = 60
 
     def __init__(self, cluster_name: str, backend: 'backends.Backend',
-                 task: 'task_lib.Task', retry_until_up: bool) -> None:
+                 task: 'task_lib.Task', max_restarts_on_errors: int,
+                 job_id: int) -> 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.
-            retry_until_up: Whether to retry until the cluster is up.
         """
         assert isinstance(backend, backends.CloudVmRayBackend), (
             'Only CloudVMRayBackend is supported.')
@@ -80,19 +63,13 @@ class StrategyExecutor:
         self.dag.add(task)
         self.cluster_name = cluster_name
         self.backend = backend
-        self.retry_until_up = retry_until_up
-
-    def __init_subclass__(cls, name: str, default: bool = False):
-        RECOVERY_STRATEGIES[name] = cls
-        if default:
-            global DEFAULT_RECOVERY_STRATEGY
-            assert DEFAULT_RECOVERY_STRATEGY is None, (
-                'Only one strategy can be default.')
-            DEFAULT_RECOVERY_STRATEGY = name
+        self.max_restarts_on_errors = max_restarts_on_errors
+        self.job_id = job_id
+        self.restart_cnt_on_failure = 0
 
     @classmethod
     def make(cls, cluster_name: str, backend: 'backends.Backend',
-             task: 'task_lib.Task', retry_until_up: bool) -> 'StrategyExecutor':
+             task: 'task_lib.Task', job_id: int) -> 'StrategyExecutor':
         """Create a strategy from a task."""
 
         resource_list = list(task.resources)
@@ -108,8 +85,19 @@ class StrategyExecutor:
         # set the new_task_resources to be the same type (list or set) as the
         # original task.resources
         task.set_resources(type(task.resources)(new_resources_list))
-        return RECOVERY_STRATEGIES[job_recovery](cluster_name, backend, task,
-                                                 retry_until_up)
+        if isinstance(job_recovery, dict):
+            job_recovery_name = job_recovery.pop(
+                'strategy', registry.JOBS_RECOVERY_STRATEGY_REGISTRY.default)
+            max_restarts_on_errors = job_recovery.pop('max_restarts_on_errors',
+                                                      0)
+        else:
+            job_recovery_name = job_recovery
+            max_restarts_on_errors = 0
+        job_recovery_strategy = (registry.JOBS_RECOVERY_STRATEGY_REGISTRY.
+                                 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)
 
     def launch(self) -> float:
         """Launch the cluster for the first time.
@@ -123,10 +111,7 @@ class StrategyExecutor:
         Raises: Please refer to the docstring of self._launch().
         """
 
-        if self.retry_until_up:
-            job_submit_at = self._launch(max_retry=None)
-        else:
-            job_submit_at = self._launch()
+        job_submit_at = self._launch(max_retry=None)
         assert job_submit_at is not None
         return job_submit_at
 
@@ -141,6 +126,8 @@ 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:
@@ -166,9 +153,9 @@ 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.
-            sky.cancel(cluster_name=self.cluster_name,
-                       all=True,
-                       _try_cancel_if_cluster_is_init=True)
+            core.cancel(cluster_name=self.cluster_name,
+                        all=True,
+                        _try_cancel_if_cluster_is_init=True)
         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.'
@@ -176,7 +163,7 @@ class StrategyExecutor:
                         f'{common_utils.format_exception(e)}\n'
                         'Terminating the cluster explicitly to ensure no '
                         'remaining job process interferes with recovery.')
-            terminate_cluster(self.cluster_name)
+            managed_job_utils.terminate_cluster(self.cluster_name)
 
     def _wait_until_job_starts_on_cluster(self) -> Optional[float]:
         """Wait for MAX_JOB_CHECKING_RETRY times until job starts on the cluster
@@ -270,8 +257,8 @@ class StrategyExecutor:
                 1. The optimizer cannot find a feasible solution.
                 2. Precheck errors: invalid cluster name, failure in getting
                 cloud user identity, or unsupported feature.
-            exceptions.SpotJobReachedMaxRetryError: This will be raised when
-                all prechecks passed but the maximum number of retries is
+            exceptions.ManagedJobReachedMaxRetriesError: This will be raised
+                when all prechecks passed but the maximum number of retries is
                 reached for `sky.launch`. The failure of `sky.launch` can be
                 due to:
                 1. Any of the underlying failover exceptions is due to resources
@@ -285,104 +272,128 @@ class StrategyExecutor:
         backoff = common_utils.Backoff(self.RETRY_INIT_GAP_SECONDS)
         while True:
             retry_cnt += 1
-            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_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}')
+            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=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.')
-
-            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 clusters '
-                            f'after {max_retry} retries.')
-                else:
+                        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}')
+                        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:
+                        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)
 
+    def should_restart_on_failure(self) -> bool:
+        """Increments counter & checks if job should be restarted on a failure.
 
-class FailoverStrategyExecutor(StrategyExecutor, name='FAILOVER',
-                               default=False):
+        Returns:
+            True if the job should be restarted, otherwise False.
+        """
+        self.restart_cnt_on_failure += 1
+        if self.restart_cnt_on_failure > self.max_restarts_on_errors:
+            return False
+        return True
+
+
+@registry.JOBS_RECOVERY_STRATEGY_REGISTRY.type_register(name='FAILOVER',
+                                                        default=False)
+class FailoverStrategyExecutor(StrategyExecutor):
     """Failover strategy: wait in same region and failover after timeout."""
 
     _MAX_RETRY_CNT = 240  # Retry for 4 hours.
 
     def __init__(self, cluster_name: str, backend: 'backends.Backend',
-                 task: 'task_lib.Task', retry_until_up: bool) -> None:
-        super().__init__(cluster_name, backend, task, retry_until_up)
+                 task: 'task_lib.Task', max_restarts_on_errors: int,
+                 job_id: int) -> None:
+        super().__init__(cluster_name, backend, task, max_restarts_on_errors,
+                         job_id)
         # 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['sky.resources.Resources'] = None
+        self._launched_resources: Optional['resources.Resources'] = None
 
     def _launch(self,
                 max_retry: Optional[int] = 3,
@@ -431,7 +442,7 @@ class FailoverStrategyExecutor(StrategyExecutor, name='FAILOVER',
             # Step 2
             logger.debug('Terminating unhealthy cluster and reset cloud '
                          'region.')
-            terminate_cluster(self.cluster_name)
+            managed_job_utils.terminate_cluster(self.cluster_name)
 
             # Step 3
             logger.debug('Relaunch the cluster  without constraining to prior '
@@ -441,23 +452,18 @@ class FailoverStrategyExecutor(StrategyExecutor, name='FAILOVER',
                                             raise_on_failure=False)
             if job_submitted_at is None:
                 # Failed to launch the cluster.
-                if self.retry_until_up:
-                    gap_seconds = self.RETRY_INIT_GAP_SECONDS
-                    logger.info('Retrying to recover the cluster in '
-                                f'{gap_seconds:.1f} seconds.')
-                    time.sleep(gap_seconds)
-                    continue
-                with ux_utils.print_exception_no_traceback():
-                    raise exceptions.ResourcesUnavailableError(
-                        f'Failed to recover the cluster after retrying '
-                        f'{self._MAX_RETRY_CNT} times.')
+                gap_seconds = self.RETRY_INIT_GAP_SECONDS
+                logger.info('Retrying to recover the cluster in '
+                            f'{gap_seconds:.1f} seconds.')
+                time.sleep(gap_seconds)
+                continue
 
             return job_submitted_at
 
 
-class EagerFailoverStrategyExecutor(FailoverStrategyExecutor,
-                                    name='EAGER_NEXT_REGION',
-                                    default=True):
+@registry.JOBS_RECOVERY_STRATEGY_REGISTRY.type_register(
+    name='EAGER_NEXT_REGION', default=True)
+class EagerFailoverStrategyExecutor(FailoverStrategyExecutor):
     """Eager failover strategy.
 
     This strategy is an extension of the FAILOVER strategy. Instead of waiting
@@ -494,7 +500,7 @@ class EagerFailoverStrategyExecutor(FailoverStrategyExecutor,
 
         # Step 1
         logger.debug('Terminating unhealthy cluster and reset cloud region.')
-        terminate_cluster(self.cluster_name)
+        managed_job_utils.terminate_cluster(self.cluster_name)
 
         # Step 2
         logger.debug('Relaunch the cluster skipping the previously launched '
@@ -529,15 +535,10 @@ class EagerFailoverStrategyExecutor(FailoverStrategyExecutor,
                                             raise_on_failure=False)
             if job_submitted_at is None:
                 # Failed to launch the cluster.
-                if self.retry_until_up:
-                    gap_seconds = self.RETRY_INIT_GAP_SECONDS
-                    logger.info('Retrying to recover the cluster in '
-                                f'{gap_seconds:.1f} seconds.')
-                    time.sleep(gap_seconds)
-                    continue
-                with ux_utils.print_exception_no_traceback():
-                    raise exceptions.ResourcesUnavailableError(
-                        f'Failed to recover the cluster after retrying '
-                        f'{self._MAX_RETRY_CNT} times.')
+                gap_seconds = self.RETRY_INIT_GAP_SECONDS
+                logger.info('Retrying to recover the cluster in '
+                            f'{gap_seconds:.1f} seconds.')
+                time.sleep(gap_seconds)
+                continue
 
             return job_submitted_at