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

sky/serve/replica_managers.py

@@ -1,6 +1,5 @@
 """ReplicaManager: handles the creation and deletion of endpoint replicas."""
 import dataclasses
-import enum
 import functools
 import multiprocessing
 from multiprocessing import pool as mp_pool
@@ -12,53 +11,56 @@ import typing
 from typing import Any, Dict, List, Optional, Tuple
 
 import colorama
-import psutil
+import filelock
 import requests
 
-import sky
 from sky import backends
 from sky import core
 from sky import exceptions
 from sky import execution
 from sky import global_user_state
 from sky import sky_logging
+from sky import task as task_lib
 from sky.backends import backend_utils
 from sky.serve import constants as serve_constants
 from sky.serve import serve_state
 from sky.serve import serve_utils
 from sky.serve import service
 from sky.serve import spot_placer
+from sky.server.requests import request_names
 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 controller_utils
 from sky.utils import env_options
+from sky.utils import resources_utils
 from sky.utils import status_lib
 from sky.utils import ux_utils
+from sky.utils import yaml_utils
 
 if typing.TYPE_CHECKING:
-    from sky import resources
     from sky.serve import service_spec
 
 logger = sky_logging.init_logger(__name__)
 
 _JOB_STATUS_FETCH_INTERVAL = 30
 _PROCESS_POOL_REFRESH_INTERVAL = 20
-# TODO(tian): Maybe let user determine this threshold
-_CONSECUTIVE_FAILURE_THRESHOLD_TIMEOUT = 180
 _RETRY_INIT_GAP_SECONDS = 60
 _DEFAULT_DRAIN_SECONDS = 120
 
-# Since sky.launch is very resource demanding, we limit the number of
-# concurrent sky.launch process to avoid overloading the machine.
-_MAX_NUM_LAUNCH = psutil.cpu_count() * 2
+# TODO(tian): Backward compatibility. Remove this after 3 minor release, i.e.
+# 0.13.0. We move the ProcessStatus to common_utils.ProcessStatus in #6666, but
+# old ReplicaInfo in database will still tries to unpickle using ProcessStatus
+# in replica_managers. We set this alias to avoid breaking changes. See #6729
+# for more details.
+ProcessStatus = common_utils.ProcessStatus
 
 
 # TODO(tian): Combine this with
 # sky/spot/recovery_strategy.py::StrategyExecutor::launch
 def launch_cluster(replica_id: int,
-                   task_yaml_path: str,
+                   service_task_yaml_path: str,
                    cluster_name: str,
                    resources_override: Optional[Dict[str, Any]] = None,
                    retry_until_up: bool = True,
@@ -78,8 +80,9 @@ def launch_cluster(replica_id: int,
                     f'{cluster_name} with resources override: '
                     f'{resources_override}')
     try:
-        config = common_utils.read_yaml(os.path.expanduser(task_yaml_path))
-        task = sky.Task.from_yaml_config(config)
+        config = yaml_utils.read_yaml(
+            os.path.expanduser(service_task_yaml_path))
+        task = task_lib.Task.from_yaml_config(config)
         if resources_override is not None:
             resources = task.resources
             overrided_resources = [
@@ -105,6 +108,8 @@ def launch_cluster(replica_id: int,
             execution.launch(task,
                              cluster_name,
                              retry_until_up=retry_until_up,
+                             _request_name=request_names.AdminPolicyRequestName.
+                             SERVE_LAUNCH_REPLICA,
                              _is_launched_by_sky_serve_controller=True)
             logger.info(f'Replica cluster {cluster_name} launched.')
         except (exceptions.InvalidClusterNameError,
@@ -173,17 +178,19 @@ def terminate_cluster(cluster_name: str,
             time.sleep(gap_seconds)
 
 
-def _get_resources_ports(task_yaml: str) -> str:
+def _get_resources_ports(service_task_yaml_path: str) -> str:
     """Get the resources ports used by the task."""
-    task = sky.Task.from_yaml(task_yaml)
+    task = task_lib.Task.from_yaml(service_task_yaml_path)
     # Already checked all ports are valid in sky.serve.core.up
     assert task.resources, task
     assert task.service is not None, task
+    if task.service.pool:
+        return '-'
     assert task.service.ports is not None, task
     return task.service.ports
 
 
-def _should_use_spot(task_yaml: str,
+def _should_use_spot(service_task_yaml_path: str,
                      resource_override: Optional[Dict[str, Any]]) -> bool:
     """Get whether the task should use spot."""
     if resource_override is not None:
@@ -191,7 +198,7 @@ def _should_use_spot(task_yaml: str,
         if use_spot_override is not None:
             assert isinstance(use_spot_override, bool)
             return use_spot_override
-    task = sky.Task.from_yaml(task_yaml)
+    task = task_lib.Task.from_yaml(service_task_yaml_path)
     spot_use_resources = [
         resources for resources in task.resources if resources.use_spot
     ]
@@ -200,6 +207,12 @@ def _should_use_spot(task_yaml: str,
     return len(spot_use_resources) == len(task.resources)
 
 
+# Every function that calls serve_state.add_or_update_replica should acquire
+# this lock. It is to prevent race condition when the replica status is updated
+# by multiple threads at the same time. The modification of replica info is
+# 2 database calls: read the whole replica info object, unpickle it, and modify
+# corresponding fields. Then it is write back to the database. We need to ensure
+# the read-modify-write operation is atomic.
 def with_lock(func):
 
     @functools.wraps(func)
@@ -210,22 +223,6 @@ def with_lock(func):
     return wrapper
 
 
-class ProcessStatus(enum.Enum):
-    """Process status."""
-
-    # The process is running
-    RUNNING = 'RUNNING'
-
-    # The process is finished and succeeded
-    SUCCEEDED = 'SUCCEEDED'
-
-    # The process is interrupted
-    INTERRUPTED = 'INTERRUPTED'
-
-    # The process failed
-    FAILED = 'FAILED'
-
-
 @dataclasses.dataclass
 class ReplicaStatusProperty:
     """Some properties that determine replica status.
@@ -237,15 +234,16 @@ class ReplicaStatusProperty:
         first_ready_time: The first time the service is ready.
         sky_down_status: Process status of sky.down.
     """
-    # None means sky.launch is not called yet.
-    sky_launch_status: Optional[ProcessStatus] = None
+    # sky.launch will always be scheduled on creation of ReplicaStatusProperty.
+    sky_launch_status: common_utils.ProcessStatus = (
+        common_utils.ProcessStatus.SCHEDULED)
     user_app_failed: bool = False
     service_ready_now: bool = False
     # None means readiness probe is not succeeded yet;
     # -1 means the initial delay seconds is exceeded.
     first_ready_time: Optional[float] = None
     # None means sky.down is not called yet.
-    sky_down_status: Optional[ProcessStatus] = None
+    sky_down_status: Optional[common_utils.ProcessStatus] = None
     # Whether the termination is caused by autoscaler's decision
     is_scale_down: bool = False
     # The replica's spot instance was preempted.
@@ -300,7 +298,7 @@ class ReplicaStatusProperty:
             (1) Job status;
             (2) Readiness probe.
         """
-        if self.sky_launch_status != ProcessStatus.SUCCEEDED:
+        if self.sky_launch_status != common_utils.ProcessStatus.SUCCEEDED:
             return False
         if self.sky_down_status is not None:
             return False
@@ -314,37 +312,43 @@ class ReplicaStatusProperty:
 
     def to_replica_status(self) -> serve_state.ReplicaStatus:
         """Convert status property to human-readable replica status."""
-        if self.sky_launch_status is None:
+        # Backward compatibility. Before we introduce ProcessStatus.SCHEDULED,
+        # we use None to represent sky.launch is not called yet.
+        if (self.sky_launch_status is None or
+                self.sky_launch_status == common_utils.ProcessStatus.SCHEDULED):
             # Pending to launch
             return serve_state.ReplicaStatus.PENDING
-        if self.sky_launch_status == ProcessStatus.RUNNING:
-            if self.sky_down_status == ProcessStatus.FAILED:
+        if self.sky_launch_status == common_utils.ProcessStatus.RUNNING:
+            if self.sky_down_status == common_utils.ProcessStatus.FAILED:
                 return serve_state.ReplicaStatus.FAILED_CLEANUP
-            if self.sky_down_status == ProcessStatus.SUCCEEDED:
+            if self.sky_down_status == common_utils.ProcessStatus.SUCCEEDED:
                 # This indicate it is a scale_down with correct teardown.
                 # Should have been cleaned from the replica table.
                 return serve_state.ReplicaStatus.UNKNOWN
             # Still launching
             return serve_state.ReplicaStatus.PROVISIONING
-        if self.sky_launch_status == ProcessStatus.INTERRUPTED:
+        if self.sky_launch_status == common_utils.ProcessStatus.INTERRUPTED:
             # sky.down is running and a scale down interrupted sky.launch
             return serve_state.ReplicaStatus.SHUTTING_DOWN
         if self.sky_down_status is not None:
             if self.preempted:
                 # Replica (spot) is preempted
                 return serve_state.ReplicaStatus.PREEMPTED
-            if self.sky_down_status == ProcessStatus.RUNNING:
+            if self.sky_down_status == common_utils.ProcessStatus.SCHEDULED:
+                # sky.down is scheduled to run, but not started yet.
+                return serve_state.ReplicaStatus.SHUTTING_DOWN
+            if self.sky_down_status == common_utils.ProcessStatus.RUNNING:
                 # sky.down is running
                 return serve_state.ReplicaStatus.SHUTTING_DOWN
-            if self.sky_launch_status == ProcessStatus.INTERRUPTED:
+            if self.sky_launch_status == common_utils.ProcessStatus.INTERRUPTED:
                 return serve_state.ReplicaStatus.SHUTTING_DOWN
-            if self.sky_down_status == ProcessStatus.FAILED:
+            if self.sky_down_status == common_utils.ProcessStatus.FAILED:
                 # sky.down failed
                 return serve_state.ReplicaStatus.FAILED_CLEANUP
             if self.user_app_failed:
                 # Failed on user setup/run
                 return serve_state.ReplicaStatus.FAILED
-            if self.sky_launch_status == ProcessStatus.FAILED:
+            if self.sky_launch_status == common_utils.ProcessStatus.FAILED:
                 # sky.launch failed
                 return serve_state.ReplicaStatus.FAILED_PROVISION
             if self.first_ready_time is None:
@@ -360,7 +364,7 @@ class ReplicaStatusProperty:
             # This indicate it is a scale_down with correct teardown.
             # Should have been cleaned from the replica table.
             return serve_state.ReplicaStatus.UNKNOWN
-        if self.sky_launch_status == ProcessStatus.FAILED:
+        if self.sky_launch_status == common_utils.ProcessStatus.FAILED:
             # sky.launch failed
             # The down process has not been started if it reaches here,
             # due to the `if self.sky_down_status is not None`` check above.
@@ -421,11 +425,12 @@ class ReplicaInfo:
                 based on the cluster name.
         """
         if cluster_record is None:
-            cluster_record = global_user_state.get_cluster_from_name(
+            handle = global_user_state.get_handle_from_cluster_name(
                 self.cluster_name)
-        if cluster_record is None:
+        else:
+            handle = cluster_record['handle']
+        if handle is None:
             return None
-        handle = cluster_record['handle']
         assert isinstance(handle, backends.CloudVmRayResourceHandle)
         return handle
 
@@ -442,10 +447,16 @@ class ReplicaInfo:
         handle = self.handle()
         if handle is None:
             return None
+        if self.replica_port == '-':
+            # This is a pool replica so there is no endpoint and it's filled
+            # with this dummy value. We return None here so that we can
+            # get the active ready replicas and perform autoscaling. Otherwise,
+            # would error out when trying to get the endpoint.
+            return None
         replica_port_int = int(self.replica_port)
         try:
-            endpoint_dict = core.endpoints(handle.cluster_name,
-                                           replica_port_int)
+            endpoint_dict = backend_utils.get_endpoints(handle.cluster_name,
+                                                        replica_port_int)
         except exceptions.ClusterNotUpError:
             return None
         endpoint = endpoint_dict.get(replica_port_int, None)
@@ -465,26 +476,36 @@ class ReplicaInfo:
                          f'replica {self.replica_id}.')
         return replica_status
 
-    def to_info_dict(self, with_handle: bool) -> Dict[str, Any]:
+    def to_info_dict(self,
+                     with_handle: bool,
+                     with_url: bool = True) -> Dict[str, Any]:
         cluster_record = global_user_state.get_cluster_from_name(
-            self.cluster_name)
+            self.cluster_name, include_user_info=False, summary_response=True)
         info_dict = {
             'replica_id': self.replica_id,
             'name': self.cluster_name,
             'status': self.status,
             'version': self.version,
-            'endpoint': self.url,
+            'endpoint': self.url if with_url else None,
             'is_spot': self.is_spot,
             'launched_at': (cluster_record['launched_at']
                             if cluster_record is not None else None),
         }
         if with_handle:
-            info_dict['handle'] = self.handle(cluster_record)
+            handle = self.handle(cluster_record)
+            info_dict['handle'] = handle
+            if handle is not None:
+                info_dict['cloud'] = repr(handle.launched_resources.cloud)
+                info_dict['region'] = handle.launched_resources.region
+                info_dict['resources_str'] = (
+                    resources_utils.get_readable_resources_repr(
+                        handle, simplified_only=True)[0])
         return info_dict
 
     def __repr__(self) -> str:
-        info_dict = self.to_info_dict(
-            with_handle=env_options.Options.SHOW_DEBUG_INFO.get())
+        show_details = env_options.Options.SHOW_DEBUG_INFO.get()
+        info_dict = self.to_info_dict(with_handle=show_details,
+                                      with_url=show_details)
         handle_str = ''
         if 'handle' in info_dict:
             handle_str = f', handle={info_dict["handle"]}'
@@ -498,6 +519,33 @@ class ReplicaInfo:
                 f'launched_at={info_dict["launched_at"]}{handle_str})')
         return info
 
+    def probe_pool(self) -> Tuple['ReplicaInfo', bool, float]:
+        """Probe the replica for pool management.
+
+        This function will check the first job status of the cluster, which is a
+        dummy job that only echoes "setup done". The success of this job means
+        the setup command is done and the replica is ready to be used. Check
+        sky/serve/server/core.py::up for more details.
+
+        Returns:
+            Tuple of (self, is_ready, probe_time).
+        """
+        probe_time = time.time()
+        try:
+            handle = backend_utils.check_cluster_available(
+                self.cluster_name, operation='probing pool')
+            if handle is None:
+                return self, False, probe_time
+            backend = backend_utils.get_backend_from_handle(handle)
+            statuses = backend.get_job_status(handle, [1], stream_logs=False)
+            if statuses[1] == job_lib.JobStatus.SUCCEEDED:
+                return self, True, probe_time
+            return self, False, probe_time
+        except Exception as e:  # pylint: disable=broad-except
+            logger.error(f'Error when probing pool of {self.cluster_name}: '
+                         f'{common_utils.format_exception(e)}.')
+            return self, False, probe_time
+
     def probe(
         self,
         readiness_path: str,
@@ -587,6 +635,7 @@ class ReplicaManager:
         self._service_name: str = service_name
         self._uptime: Optional[float] = None
         self._update_mode = serve_utils.DEFAULT_UPDATE_MODE
+        self._is_pool: bool = spec.pool
         header_keys = None
         if spec.readiness_headers is not None:
             header_keys = list(spec.readiness_headers.keys())
@@ -600,6 +649,15 @@ class ReplicaManager:
         # Oldest version among the currently provisioned and launched replicas
         self.least_recent_version: int = serve_constants.INITIAL_VERSION
 
+    def _consecutive_failure_threshold_timeout(self) -> int:
+        """The timeout for the consecutive failure threshold in seconds.
+
+        We reduce the timeout for pool to 10 seconds to make the pool more
+        responsive to the failure.
+        """
+        # TODO(tian): Maybe let user determine this threshold
+        return 10 if self._is_pool else 180
+
     def scale_up(self,
                  resources_override: Optional[Dict[str, Any]] = None) -> None:
         """Scale up the service by 1 replica with resources_override.
@@ -634,10 +692,10 @@ class SkyPilotReplicaManager(ReplicaManager):
     """
 
     def __init__(self, service_name: str, spec: 'service_spec.SkyServiceSpec',
-                 task_yaml_path: str) -> None:
+                 service_task_yaml_path: str) -> None:
         super().__init__(service_name, spec)
-        self._task_yaml_path = task_yaml_path
-        task = sky.Task.from_yaml(task_yaml_path)
+        self.service_task_yaml_path = service_task_yaml_path
+        task = task_lib.Task.from_yaml(service_task_yaml_path)
         self._spot_placer: Optional[spot_placer.SpotPlacer] = (
             spot_placer.SpotPlacer.from_task(spec, task))
         # TODO(tian): Store launch/down pid in the replica table, to make the
@@ -657,6 +715,7 @@ class SkyPilotReplicaManager(ReplicaManager):
 
         self._recover_replica_operations()
 
+    @with_lock
     def _recover_replica_operations(self):
         """Let's see are there something to do for ReplicaManager in a
         recovery run"""
@@ -697,9 +756,8 @@ class SkyPilotReplicaManager(ReplicaManager):
     # Replica management functions #
     ################################
 
-    # Adding lock here to make sure spot placer's current locations are
-    # consistent with the replicas' status.
-    @with_lock
+    # We don't need to add lock here since every caller of this function
+    # will acquire the lock.
     def _launch_replica(
         self,
         replica_id: int,
@@ -714,7 +772,8 @@ class SkyPilotReplicaManager(ReplicaManager):
             self._service_name, replica_id)
         log_file_name = serve_utils.generate_replica_launch_log_file_name(
             self._service_name, replica_id)
-        use_spot = _should_use_spot(self._task_yaml_path, resources_override)
+        use_spot = _should_use_spot(self.service_task_yaml_path,
+                                    resources_override)
         retry_until_up = True
         location = None
         if use_spot and self._spot_placer is not None:
@@ -742,10 +801,10 @@ class SkyPilotReplicaManager(ReplicaManager):
                 launch_cluster,
                 log_file_name,
             ).run,
-            args=(replica_id, self._task_yaml_path, cluster_name,
+            args=(replica_id, self.service_task_yaml_path, cluster_name,
                   resources_override, retry_until_up),
         )
-        replica_port = _get_resources_ports(self._task_yaml_path)
+        replica_port = _get_resources_ports(self.service_task_yaml_path)
 
         info = ReplicaInfo(replica_id, cluster_name, replica_port, use_spot,
                            location, self.latest_version, resources_override)
@@ -754,11 +813,61 @@ class SkyPilotReplicaManager(ReplicaManager):
         # to avoid too many sky.launch running at the same time.
         self._launch_process_pool[replica_id] = p
 
+    @with_lock
     def scale_up(self,
                  resources_override: Optional[Dict[str, Any]] = None) -> None:
         self._launch_replica(self._next_replica_id, resources_override)
         self._next_replica_id += 1
 
+    def _handle_sky_down_finish(self, info: ReplicaInfo, exitcode: int) -> None:
+        if exitcode != 0:
+            logger.error(f'Down process for replica {info.replica_id} '
+                         f'exited abnormally with code {exitcode}.')
+            info.status_property.sky_down_status = (
+                common_utils.ProcessStatus.FAILED)
+        else:
+            info.status_property.sky_down_status = (
+                common_utils.ProcessStatus.SUCCEEDED)
+        # Failed replica still count as a replica. In our current design, we
+        # want to fail early if user code have any error. This will prevent
+        # infinite loop of teardown and re-provision. However, there is a
+        # special case that if the replica is UP for longer than
+        # initial_delay_seconds, we assume it is just some random failure and
+        # we should restart the replica. Please refer to the implementation of
+        # `is_scale_down_succeeded` for more details.
+        # TODO(tian): Currently, restart replicas that failed within
+        # initial_delay_seconds is not supported. We should add it
+        # later when we support `sky serve update`.
+        removal_reason = None
+        if info.status_property.is_scale_down:
+            # This means the cluster is deleted due to an autoscaler
+            # decision or the cluster is recovering from preemption.
+            # Delete the replica info so it won't count as a replica.
+            if info.status_property.preempted:
+                removal_reason = 'for preemption recovery'
+            else:
+                removal_reason = 'normally'
+        # Don't keep failed record for version mismatch replicas,
+        # since user should fixed the error before update.
+        elif info.version != self.latest_version:
+            removal_reason = 'for version outdated'
+        elif info.status_property.purged:
+            removal_reason = 'for purge'
+        elif info.status_property.failed_spot_availability:
+            removal_reason = 'for spot availability failure'
+        else:
+            logger.info(f'Termination of replica {info.replica_id} '
+                        'finished. Replica info is kept since some '
+                        'failure detected.')
+            serve_state.add_or_update_replica(self._service_name,
+                                              info.replica_id, info)
+        if removal_reason is not None:
+            serve_state.remove_replica(self._service_name, info.replica_id)
+            logger.info(f'Replica {info.replica_id} removed from the '
+                        f'replica table {removal_reason}.')
+
+    # We don't need to add lock here since every caller of this function
+    # will acquire the lock.
     def _terminate_replica(self,
                            replica_id: int,
                            sync_down_logs: bool,
@@ -776,7 +885,8 @@ class SkyPilotReplicaManager(ReplicaManager):
             info = serve_state.get_replica_info_from_id(self._service_name,
                                                         replica_id)
             assert info is not None
-            info.status_property.sky_launch_status = ProcessStatus.INTERRUPTED
+            info.status_property.sky_launch_status = (
+                common_utils.ProcessStatus.INTERRUPTED)
             serve_state.add_or_update_replica(self._service_name, replica_id,
                                               info)
             launch_process = self._launch_process_pool[replica_id]
@@ -820,9 +930,9 @@ class SkyPilotReplicaManager(ReplicaManager):
             assert isinstance(handle, backends.CloudVmRayResourceHandle)
             replica_job_logs_dir = os.path.join(constants.SKY_LOGS_DIRECTORY,
                                                 'replica_jobs')
-            job_log_file_name = (
-                controller_utils.download_and_stream_latest_job_log(
-                    backend, handle, replica_job_logs_dir))
+            job_ids = ['1'] if self._is_pool else None
+            job_log_file_name = controller_utils.download_and_stream_job_log(
+                backend, handle, replica_job_logs_dir, job_ids)
             if job_log_file_name is not None:
                 logger.info(f'\n== End of logs (Replica: {replica_id}) ==')
                 with open(log_file_name, 'a',
@@ -848,18 +958,30 @@ class SkyPilotReplicaManager(ReplicaManager):
 
         logger.info(f'preempted: {info.status_property.preempted}, '
                     f'replica_id: {replica_id}')
+        info.status_property.is_scale_down = is_scale_down
+        info.status_property.purged = purge
+
+        # If the cluster does not exist, it means either the cluster never
+        # exists (e.g., the cluster is scaled down before it gets a chance to
+        # provision) or the cluster is preempted and cleaned up by the status
+        # refresh. In this case, we skip spawning a new down process to save
+        # controller resources.
+        if not global_user_state.cluster_with_name_exists(info.cluster_name):
+            self._handle_sky_down_finish(info, exitcode=0)
+            return
+
+        # Otherwise, start the process to terminate the cluster.
         p = multiprocessing.Process(
             target=ux_utils.RedirectOutputForProcess(terminate_cluster,
                                                      log_file_name, 'a').run,
             args=(info.cluster_name, replica_drain_delay_seconds),
         )
-        info.status_property.sky_down_status = ProcessStatus.RUNNING
-        info.status_property.is_scale_down = is_scale_down
-        info.status_property.purged = purge
+        info.status_property.sky_down_status = (
+            common_utils.ProcessStatus.SCHEDULED)
         serve_state.add_or_update_replica(self._service_name, replica_id, info)
-        p.start()
         self._down_process_pool[replica_id] = p
 
+    @with_lock
     def scale_down(self, replica_id: int, purge: bool = False) -> None:
         self._terminate_replica(
             replica_id,
@@ -868,6 +990,8 @@ class SkyPilotReplicaManager(ReplicaManager):
             is_scale_down=True,
             purge=purge)
 
+    # We don't need to add lock here since every caller of this function
+    # will acquire the lock.
     def _handle_preemption(self, info: ReplicaInfo) -> bool:
         """Handle preemption of the replica if any error happened.
 
@@ -930,18 +1054,19 @@ class SkyPilotReplicaManager(ReplicaManager):
         # To avoid `dictionary changed size during iteration` error.
         launch_process_pool_snapshot = list(self._launch_process_pool.items())
         for replica_id, p in launch_process_pool_snapshot:
-            if not p.is_alive():
+            if p.is_alive():
+                continue
+            with filelock.FileLock(controller_utils.get_resources_lock_path()):
                 info = serve_state.get_replica_info_from_id(
                     self._service_name, replica_id)
                 assert info is not None, replica_id
                 error_in_sky_launch = False
                 if info.status == serve_state.ReplicaStatus.PENDING:
                     # sky.launch not started yet
-                    if (serve_state.total_number_provisioning_replicas() <
-                            _MAX_NUM_LAUNCH):
+                    if controller_utils.can_provision():
                         p.start()
                         info.status_property.sky_launch_status = (
-                            ProcessStatus.RUNNING)
+                            common_utils.ProcessStatus.RUNNING)
                 else:
                     # sky.launch finished
                     # TODO(tian): Try-catch in process, and have an enum return
@@ -958,11 +1083,11 @@ class SkyPilotReplicaManager(ReplicaManager):
                             f'exited abnormally with code {p.exitcode}.'
                             ' Terminating...')
                         info.status_property.sky_launch_status = (
-                            ProcessStatus.FAILED)
+                            common_utils.ProcessStatus.FAILED)
                         error_in_sky_launch = True
                     else:
                         info.status_property.sky_launch_status = (
-                            ProcessStatus.SUCCEEDED)
+                            common_utils.ProcessStatus.SUCCEEDED)
                     if self._spot_placer is not None and info.is_spot:
                         # TODO(tian): Currently, we set the location to
                         # preemptive if the launch process failed. This is
@@ -990,59 +1115,25 @@ class SkyPilotReplicaManager(ReplicaManager):
                                             replica_drain_delay_seconds=0)
         down_process_pool_snapshot = list(self._down_process_pool.items())
         for replica_id, p in down_process_pool_snapshot:
-            if not p.is_alive():
-                logger.info(
-                    f'Terminate process for replica {replica_id} finished.')
-                del self._down_process_pool[replica_id]
-                info = serve_state.get_replica_info_from_id(
-                    self._service_name, replica_id)
-                assert info is not None, replica_id
-                if p.exitcode != 0:
-                    logger.error(f'Down process for replica {replica_id} '
-                                 f'exited abnormally with code {p.exitcode}.')
-                    info.status_property.sky_down_status = (
-                        ProcessStatus.FAILED)
-                else:
+            if p.is_alive():
+                continue
+            info = serve_state.get_replica_info_from_id(self._service_name,
+                                                        replica_id)
+            assert info is not None, replica_id
+            if (info.status_property.sky_down_status ==
+                    common_utils.ProcessStatus.SCHEDULED):
+                # sky.down not started yet
+                if controller_utils.can_terminate():
+                    p.start()
                     info.status_property.sky_down_status = (
-                        ProcessStatus.SUCCEEDED)
-                # Failed replica still count as a replica. In our current
-                # design, we want to fail early if user code have any error.
-                # This will prevent infinite loop of teardown and
-                # re-provision. However, there is a special case that if the
-                # replica is UP for longer than initial_delay_seconds, we
-                # assume it is just some random failure and we should restart
-                # the replica. Please refer to the implementation of
-                # `is_scale_down_succeeded` for more details.
-                # TODO(tian): Currently, restart replicas that failed within
-                # initial_delay_seconds is not supported. We should add it
-                # later when we support `sky serve update`.
-                removal_reason = None
-                if info.status_property.is_scale_down:
-                    # This means the cluster is deleted due to an autoscaler
-                    # decision or the cluster is recovering from preemption.
-                    # Delete the replica info so it won't count as a replica.
-                    if info.status_property.preempted:
-                        removal_reason = 'for preemption recovery'
-                    else:
-                        removal_reason = 'normally'
-                # Don't keep failed record for version mismatch replicas,
-                # since user should fixed the error before update.
-                elif info.version != self.latest_version:
-                    removal_reason = 'for version outdated'
-                elif info.status_property.purged:
-                    removal_reason = 'for purge'
-                elif info.status_property.failed_spot_availability:
-                    removal_reason = 'for spot availability failure'
-                else:
-                    logger.info(f'Termination of replica {replica_id} '
-                                'finished. Replica info is kept since some '
-                                'failure detected.')
+                        common_utils.ProcessStatus.RUNNING)
                     serve_state.add_or_update_replica(self._service_name,
                                                       replica_id, info)
-                if removal_reason is not None:
-                    serve_state.remove_replica(self._service_name, replica_id)
-                    logger.info(f'Replica {replica_id} removed from the '
-                                f'replica table {removal_reason}.')
+            else:
+                logger.info(
+                    f'Terminate process for replica {replica_id} finished.')
+                del self._down_process_pool[replica_id]
+                self._handle_sky_down_finish(info, exitcode=p.exitcode)
 
         # Clean old version
         replica_infos = serve_state.get_replica_infos(self._service_name)
@@ -1098,9 +1189,10 @@ class SkyPilotReplicaManager(ReplicaManager):
             handle = info.handle()
             assert handle is not None, info
             # Use None to fetch latest job, which stands for user task job
+            job_ids = [1] if self._is_pool else None
             try:
                 job_statuses = backend.get_job_status(handle,
-                                                      None,
+                                                      job_ids,
                                                       stream_logs=False)
             except exceptions.CommandError:
                 # If the job status fetch failed, it is likely that the
@@ -1110,7 +1202,8 @@ class SkyPilotReplicaManager(ReplicaManager):
                     continue
                 # Re-raise the exception if it is not preempted.
                 raise
-            job_status = list(job_statuses.values())[0]
+            job_status = job_statuses[1] if self._is_pool else list(
+                job_statuses.values())[0]
             if job_status in job_lib.JobStatus.user_code_failure_states():
                 info.status_property.user_app_failed = True
                 serve_state.add_or_update_replica(self._service_name,
@@ -1154,18 +1247,24 @@ class SkyPilotReplicaManager(ReplicaManager):
             for info in infos:
                 if not info.status_property.should_track_service_status():
                     continue
-                replica_to_probe.append(
-                    f'replica_{info.replica_id}(url={info.url})')
-                probe_futures.append(
-                    pool.apply_async(
-                        info.probe,
-                        (
-                            self._get_readiness_path(info.version),
-                            self._get_post_data(info.version),
-                            self._get_readiness_timeout_seconds(info.version),
-                            self._get_readiness_headers(info.version),
-                        ),
-                    ),)
+                if self._is_pool:
+                    replica_to_probe.append(f'replica_{info.replica_id}(cluster'
+                                            f'_name={info.cluster_name})')
+                    probe_futures.append(pool.apply_async(info.probe_pool))
+                else:
+                    replica_to_probe.append(
+                        f'replica_{info.replica_id}(url={info.url})')
+                    probe_futures.append(
+                        pool.apply_async(
+                            info.probe,
+                            (
+                                self._get_readiness_path(info.version),
+                                self._get_post_data(info.version),
+                                self._get_readiness_timeout_seconds(
+                                    info.version),
+                                self._get_readiness_headers(info.version),
+                            ),
+                        ),)
             logger.info(f'Replicas to probe: {", ".join(replica_to_probe)}')
 
             # Since futures.as_completed will return futures in the order of
@@ -1202,8 +1301,9 @@ class SkyPilotReplicaManager(ReplicaManager):
                         consecutive_failure_time = (
                             info.consecutive_failure_times[-1] -
                             info.consecutive_failure_times[0])
-                        if (consecutive_failure_time >=
-                                _CONSECUTIVE_FAILURE_THRESHOLD_TIMEOUT):
+                        failure_threshold = (
+                            self._consecutive_failure_threshold_timeout())
+                        if consecutive_failure_time >= failure_threshold:
                             logger.info(
                                 f'Replica {info.replica_id} is not ready for '
                                 'too long and exceeding consecutive failure '
@@ -1214,8 +1314,7 @@ class SkyPilotReplicaManager(ReplicaManager):
                                 f'Replica {info.replica_id} is not ready '
                                 'but within consecutive failure threshold '
                                 f'({consecutive_failure_time}s / '
-                                f'{_CONSECUTIVE_FAILURE_THRESHOLD_TIMEOUT}s). '
-                                'Skipping.')
+                                f'{failure_threshold}s). Skipping.')
                     else:
                         initial_delay_seconds = self._get_initial_delay_seconds(
                             info.version)
@@ -1290,11 +1389,11 @@ class SkyPilotReplicaManager(ReplicaManager):
             logger.error(f'Invalid version: {version}, '
                          f'latest version: {self.latest_version}')
             return
-        task_yaml_path = serve_utils.generate_task_yaml_file_name(
+        service_task_yaml_path = serve_utils.generate_task_yaml_file_name(
             self._service_name, version)
         serve_state.add_or_update_version(self._service_name, version, spec)
         self.latest_version = version
-        self._task_yaml_path = task_yaml_path
+        self.service_task_yaml_path = service_task_yaml_path
         self._update_mode = update_mode
 
         # Reuse all replicas that have the same config as the new version
@@ -1302,32 +1401,40 @@ class SkyPilotReplicaManager(ReplicaManager):
         # the latest version. This can significantly improve the speed
         # for updating an existing service with only config changes to the
         # service specs, e.g. scale down the service.
-        new_config = common_utils.read_yaml(os.path.expanduser(task_yaml_path))
+        new_config = yaml_utils.read_yaml(
+            os.path.expanduser(service_task_yaml_path))
         # Always create new replicas and scale down old ones when file_mounts
         # are not empty.
         if new_config.get('file_mounts', None) != {}:
             return
-        for key in ['service']:
-            new_config.pop(key)
+        for key in ['service', 'pool', '_user_specified_yaml']:
+            new_config.pop(key, None)
+        new_config_any_of = new_config.get('resources', {}).pop('any_of', [])
+
         replica_infos = serve_state.get_replica_infos(self._service_name)
         for info in replica_infos:
             if info.version < version and not info.is_terminal:
                 # Assume user does not change the yaml file on the controller.
-                old_task_yaml_path = serve_utils.generate_task_yaml_file_name(
-                    self._service_name, info.version)
-                old_config = common_utils.read_yaml(
-                    os.path.expanduser(old_task_yaml_path))
-                for key in ['service']:
-                    old_config.pop(key)
+                old_service_task_yaml_path = (
+                    serve_utils.generate_task_yaml_file_name(
+                        self._service_name, info.version))
+                old_config = yaml_utils.read_yaml(
+                    os.path.expanduser(old_service_task_yaml_path))
+                for key in ['service', 'pool', '_user_specified_yaml']:
+                    old_config.pop(key, None)
                 # Bump replica version if all fields except for service are
                 # the same.
                 # Here, we manually convert the any_of field to a set to avoid
                 # only the difference in the random order of the any_of fields.
                 old_config_any_of = old_config.get('resources',
                                                    {}).pop('any_of', [])
-                new_config_any_of = new_config.get('resources',
-                                                   {}).pop('any_of', [])
-                if set(old_config_any_of) != set(new_config_any_of):
+
+                if (resources_utils.normalize_any_of_resources_config(
+                        old_config_any_of) != resources_utils.
+                        normalize_any_of_resources_config(new_config_any_of)):
+                    logger.info('Replica config changed (any_of), skipping. '
+                                f'old: {old_config_any_of}, '
+                                f'new: {new_config_any_of}')
                     continue
                 # File mounts should both be empty, as update always
                 # create new buckets if they are not empty.
@@ -1341,6 +1448,10 @@ class SkyPilotReplicaManager(ReplicaManager):
                     info.version = version
                     serve_state.add_or_update_replica(self._service_name,
                                                       info.replica_id, info)
+                else:
+                    logger.info('Replica config changed (rest), skipping. '
+                                f'old: {old_config}, '
+                                f'new: {new_config}')
 
     def _get_version_spec(self, version: int) -> 'service_spec.SkyServiceSpec':
         spec = serve_state.get_spec(self._service_name, version)