skypilot-nightly 1.0.0.dev20250905__py3-none-any.whl → 1.0.0.dev20251203__py3-none-any.whl

sky/schemas/generated/servev1_pb2.pyi

@@ -0,0 +1,115 @@
+from google.protobuf.internal import containers as _containers
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from typing import ClassVar as _ClassVar, Iterable as _Iterable, Mapping as _Mapping, Optional as _Optional, Union as _Union
+
+DESCRIPTOR: _descriptor.FileDescriptor
+
+class ServiceNames(_message.Message):
+    __slots__ = ("names",)
+    NAMES_FIELD_NUMBER: _ClassVar[int]
+    names: _containers.RepeatedScalarFieldContainer[str]
+    def __init__(self, names: _Optional[_Iterable[str]] = ...) -> None: ...
+
+class ServiceStatus(_message.Message):
+    __slots__ = ("status",)
+    class StatusEntry(_message.Message):
+        __slots__ = ("key", "value")
+        KEY_FIELD_NUMBER: _ClassVar[int]
+        VALUE_FIELD_NUMBER: _ClassVar[int]
+        key: str
+        value: str
+        def __init__(self, key: _Optional[str] = ..., value: _Optional[str] = ...) -> None: ...
+    STATUS_FIELD_NUMBER: _ClassVar[int]
+    status: _containers.ScalarMap[str, str]
+    def __init__(self, status: _Optional[_Mapping[str, str]] = ...) -> None: ...
+
+class GetServiceStatusRequest(_message.Message):
+    __slots__ = ("service_names", "pool")
+    SERVICE_NAMES_FIELD_NUMBER: _ClassVar[int]
+    POOL_FIELD_NUMBER: _ClassVar[int]
+    service_names: ServiceNames
+    pool: bool
+    def __init__(self, service_names: _Optional[_Union[ServiceNames, _Mapping]] = ..., pool: bool = ...) -> None: ...
+
+class GetServiceStatusResponse(_message.Message):
+    __slots__ = ("statuses",)
+    STATUSES_FIELD_NUMBER: _ClassVar[int]
+    statuses: _containers.RepeatedCompositeFieldContainer[ServiceStatus]
+    def __init__(self, statuses: _Optional[_Iterable[_Union[ServiceStatus, _Mapping]]] = ...) -> None: ...
+
+class AddVersionRequest(_message.Message):
+    __slots__ = ("service_name",)
+    SERVICE_NAME_FIELD_NUMBER: _ClassVar[int]
+    service_name: str
+    def __init__(self, service_name: _Optional[str] = ...) -> None: ...
+
+class AddVersionResponse(_message.Message):
+    __slots__ = ("version",)
+    VERSION_FIELD_NUMBER: _ClassVar[int]
+    version: int
+    def __init__(self, version: _Optional[int] = ...) -> None: ...
+
+class TerminateServicesRequest(_message.Message):
+    __slots__ = ("service_names", "purge", "pool")
+    SERVICE_NAMES_FIELD_NUMBER: _ClassVar[int]
+    PURGE_FIELD_NUMBER: _ClassVar[int]
+    POOL_FIELD_NUMBER: _ClassVar[int]
+    service_names: ServiceNames
+    purge: bool
+    pool: bool
+    def __init__(self, service_names: _Optional[_Union[ServiceNames, _Mapping]] = ..., purge: bool = ..., pool: bool = ...) -> None: ...
+
+class TerminateServicesResponse(_message.Message):
+    __slots__ = ("message",)
+    MESSAGE_FIELD_NUMBER: _ClassVar[int]
+    message: str
+    def __init__(self, message: _Optional[str] = ...) -> None: ...
+
+class TerminateReplicaRequest(_message.Message):
+    __slots__ = ("service_name", "replica_id", "purge")
+    SERVICE_NAME_FIELD_NUMBER: _ClassVar[int]
+    REPLICA_ID_FIELD_NUMBER: _ClassVar[int]
+    PURGE_FIELD_NUMBER: _ClassVar[int]
+    service_name: str
+    replica_id: int
+    purge: bool
+    def __init__(self, service_name: _Optional[str] = ..., replica_id: _Optional[int] = ..., purge: bool = ...) -> None: ...
+
+class TerminateReplicaResponse(_message.Message):
+    __slots__ = ("message",)
+    MESSAGE_FIELD_NUMBER: _ClassVar[int]
+    message: str
+    def __init__(self, message: _Optional[str] = ...) -> None: ...
+
+class WaitServiceRegistrationRequest(_message.Message):
+    __slots__ = ("service_name", "job_id", "pool")
+    SERVICE_NAME_FIELD_NUMBER: _ClassVar[int]
+    JOB_ID_FIELD_NUMBER: _ClassVar[int]
+    POOL_FIELD_NUMBER: _ClassVar[int]
+    service_name: str
+    job_id: int
+    pool: bool
+    def __init__(self, service_name: _Optional[str] = ..., job_id: _Optional[int] = ..., pool: bool = ...) -> None: ...
+
+class WaitServiceRegistrationResponse(_message.Message):
+    __slots__ = ("lb_port",)
+    LB_PORT_FIELD_NUMBER: _ClassVar[int]
+    lb_port: int
+    def __init__(self, lb_port: _Optional[int] = ...) -> None: ...
+
+class UpdateServiceRequest(_message.Message):
+    __slots__ = ("service_name", "version", "mode", "pool")
+    SERVICE_NAME_FIELD_NUMBER: _ClassVar[int]
+    VERSION_FIELD_NUMBER: _ClassVar[int]
+    MODE_FIELD_NUMBER: _ClassVar[int]
+    POOL_FIELD_NUMBER: _ClassVar[int]
+    service_name: str
+    version: int
+    mode: str
+    pool: bool
+    def __init__(self, service_name: _Optional[str] = ..., version: _Optional[int] = ..., mode: _Optional[str] = ..., pool: bool = ...) -> None: ...
+
+class UpdateServiceResponse(_message.Message):
+    __slots__ = ()
+    def __init__(self) -> None: ...

sky/schemas/generated/servev1_pb2_grpc.py

@@ -0,0 +1,322 @@
+# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+"""Client and server classes corresponding to protobuf-defined services."""
+import grpc
+import warnings
+
+from sky.schemas.generated import servev1_pb2 as sky_dot_schemas_dot_generated_dot_servev1__pb2
+
+GRPC_GENERATED_VERSION = '1.63.0'
+GRPC_VERSION = grpc.__version__
+EXPECTED_ERROR_RELEASE = '1.65.0'
+SCHEDULED_RELEASE_DATE = 'June 25, 2024'
+_version_not_supported = False
+
+try:
+    from grpc._utilities import first_version_is_lower
+    _version_not_supported = first_version_is_lower(GRPC_VERSION, GRPC_GENERATED_VERSION)
+except ImportError:
+    _version_not_supported = True
+
+if _version_not_supported:
+    warnings.warn(
+        f'The grpc package installed is at version {GRPC_VERSION},'
+        + f' but the generated code in sky/schemas/generated/servev1_pb2_grpc.py depends on'
+        + f' grpcio>={GRPC_GENERATED_VERSION}.'
+        + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}'
+        + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.'
+        + f' This warning will become an error in {EXPECTED_ERROR_RELEASE},'
+        + f' scheduled for release on {SCHEDULED_RELEASE_DATE}.',
+        RuntimeWarning
+    )
+
+
+class ServeServiceStub(object):
+    """Missing associated documentation comment in .proto file."""
+
+    def __init__(self, channel):
+        """Constructor.
+
+        Args:
+            channel: A grpc.Channel.
+        """
+        self.GetServiceStatus = channel.unary_unary(
+                '/serve.v1.ServeService/GetServiceStatus',
+                request_serializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.GetServiceStatusRequest.SerializeToString,
+                response_deserializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.GetServiceStatusResponse.FromString,
+                _registered_method=True)
+        self.AddVersion = channel.unary_unary(
+                '/serve.v1.ServeService/AddVersion',
+                request_serializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.AddVersionRequest.SerializeToString,
+                response_deserializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.AddVersionResponse.FromString,
+                _registered_method=True)
+        self.TerminateServices = channel.unary_unary(
+                '/serve.v1.ServeService/TerminateServices',
+                request_serializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.TerminateServicesRequest.SerializeToString,
+                response_deserializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.TerminateServicesResponse.FromString,
+                _registered_method=True)
+        self.TerminateReplica = channel.unary_unary(
+                '/serve.v1.ServeService/TerminateReplica',
+                request_serializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.TerminateReplicaRequest.SerializeToString,
+                response_deserializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.TerminateReplicaResponse.FromString,
+                _registered_method=True)
+        self.WaitServiceRegistration = channel.unary_unary(
+                '/serve.v1.ServeService/WaitServiceRegistration',
+                request_serializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.WaitServiceRegistrationRequest.SerializeToString,
+                response_deserializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.WaitServiceRegistrationResponse.FromString,
+                _registered_method=True)
+        self.UpdateService = channel.unary_unary(
+                '/serve.v1.ServeService/UpdateService',
+                request_serializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.UpdateServiceRequest.SerializeToString,
+                response_deserializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.UpdateServiceResponse.FromString,
+                _registered_method=True)
+
+
+class ServeServiceServicer(object):
+    """Missing associated documentation comment in .proto file."""
+
+    def GetServiceStatus(self, request, context):
+        """Get status of service.
+        """
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def AddVersion(self, request, context):
+        """Add version to service.
+        """
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def TerminateServices(self, request, context):
+        """Terminate services.
+        """
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def TerminateReplica(self, request, context):
+        """Terminate replica.
+        """
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def WaitServiceRegistration(self, request, context):
+        """Wait for service registration.
+        """
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def UpdateService(self, request, context):
+        """Update service.
+        """
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+
+def add_ServeServiceServicer_to_server(servicer, server):
+    rpc_method_handlers = {
+            'GetServiceStatus': grpc.unary_unary_rpc_method_handler(
+                    servicer.GetServiceStatus,
+                    request_deserializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.GetServiceStatusRequest.FromString,
+                    response_serializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.GetServiceStatusResponse.SerializeToString,
+            ),
+            'AddVersion': grpc.unary_unary_rpc_method_handler(
+                    servicer.AddVersion,
+                    request_deserializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.AddVersionRequest.FromString,
+                    response_serializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.AddVersionResponse.SerializeToString,
+            ),
+            'TerminateServices': grpc.unary_unary_rpc_method_handler(
+                    servicer.TerminateServices,
+                    request_deserializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.TerminateServicesRequest.FromString,
+                    response_serializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.TerminateServicesResponse.SerializeToString,
+            ),
+            'TerminateReplica': grpc.unary_unary_rpc_method_handler(
+                    servicer.TerminateReplica,
+                    request_deserializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.TerminateReplicaRequest.FromString,
+                    response_serializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.TerminateReplicaResponse.SerializeToString,
+            ),
+            'WaitServiceRegistration': grpc.unary_unary_rpc_method_handler(
+                    servicer.WaitServiceRegistration,
+                    request_deserializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.WaitServiceRegistrationRequest.FromString,
+                    response_serializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.WaitServiceRegistrationResponse.SerializeToString,
+            ),
+            'UpdateService': grpc.unary_unary_rpc_method_handler(
+                    servicer.UpdateService,
+                    request_deserializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.UpdateServiceRequest.FromString,
+                    response_serializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.UpdateServiceResponse.SerializeToString,
+            ),
+    }
+    generic_handler = grpc.method_handlers_generic_handler(
+            'serve.v1.ServeService', rpc_method_handlers)
+    server.add_generic_rpc_handlers((generic_handler,))
+
+
+ # This class is part of an EXPERIMENTAL API.
+class ServeService(object):
+    """Missing associated documentation comment in .proto file."""
+
+    @staticmethod
+    def GetServiceStatus(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            '/serve.v1.ServeService/GetServiceStatus',
+            sky_dot_schemas_dot_generated_dot_servev1__pb2.GetServiceStatusRequest.SerializeToString,
+            sky_dot_schemas_dot_generated_dot_servev1__pb2.GetServiceStatusResponse.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)
+
+    @staticmethod
+    def AddVersion(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            '/serve.v1.ServeService/AddVersion',
+            sky_dot_schemas_dot_generated_dot_servev1__pb2.AddVersionRequest.SerializeToString,
+            sky_dot_schemas_dot_generated_dot_servev1__pb2.AddVersionResponse.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)
+
+    @staticmethod
+    def TerminateServices(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            '/serve.v1.ServeService/TerminateServices',
+            sky_dot_schemas_dot_generated_dot_servev1__pb2.TerminateServicesRequest.SerializeToString,
+            sky_dot_schemas_dot_generated_dot_servev1__pb2.TerminateServicesResponse.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)
+
+    @staticmethod
+    def TerminateReplica(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            '/serve.v1.ServeService/TerminateReplica',
+            sky_dot_schemas_dot_generated_dot_servev1__pb2.TerminateReplicaRequest.SerializeToString,
+            sky_dot_schemas_dot_generated_dot_servev1__pb2.TerminateReplicaResponse.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)
+
+    @staticmethod
+    def WaitServiceRegistration(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            '/serve.v1.ServeService/WaitServiceRegistration',
+            sky_dot_schemas_dot_generated_dot_servev1__pb2.WaitServiceRegistrationRequest.SerializeToString,
+            sky_dot_schemas_dot_generated_dot_servev1__pb2.WaitServiceRegistrationResponse.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)
+
+    @staticmethod
+    def UpdateService(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            '/serve.v1.ServeService/UpdateService',
+            sky_dot_schemas_dot_generated_dot_servev1__pb2.UpdateServiceRequest.SerializeToString,
+            sky_dot_schemas_dot_generated_dot_servev1__pb2.UpdateServiceResponse.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)

sky/serve/autoscalers.py

@@ -411,6 +411,8 @@ class _AutoscalerWithHysteresis(Autoscaler):
         # `_set_target_num_replicas_with_hysteresis` to have the replicas
         # quickly scale after each update.
         self.target_num_replicas = self._calculate_target_num_replicas()
+        logger.debug(f'Target number of replicas: {self.target_num_replicas}'
+                     'after update_version.')
         # Cleanup hysteresis counters.
         self.upscale_counter = 0
         self.downscale_counter = 0

sky/serve/client/impl.py

@@ -8,6 +8,7 @@ import click
 from sky.client import common as client_common
 from sky.server import common as server_common
 from sky.server.requests import payloads
+from sky.server.requests import request_names
 from sky.utils import admin_policy_utils
 from sky.utils import dag_utils
 
@@ -32,7 +33,9 @@ def up(
 
     dag = dag_utils.convert_entrypoint_to_dag(task)
     with admin_policy_utils.apply_and_use_config_in_current_request(
-            dag, at_client_side=True) as dag:
+            dag,
+            request_name=request_names.AdminPolicyRequestName.SERVE_UP,
+            at_client_side=True) as dag:
         sdk.validate(dag)
         request_id = sdk.optimize(dag)
         sdk.stream_and_get(request_id)
@@ -77,7 +80,9 @@ def update(
 
     dag = dag_utils.convert_entrypoint_to_dag(task)
     with admin_policy_utils.apply_and_use_config_in_current_request(
-            dag, at_client_side=True) as dag:
+            dag,
+            request_name=request_names.AdminPolicyRequestName.SERVE_UPDATE,
+            at_client_side=True) as dag:
         sdk.validate(dag)
         request_id = sdk.optimize(dag)
         sdk.stream_and_get(request_id)
@@ -105,7 +110,8 @@ def update(
 
 
 def apply(
-    task: Union['sky.Task', 'sky.Dag'],
+    task: Optional[Union['sky.Task', 'sky.Dag']],
+    workers: Optional[int],
     service_name: str,
     mode: 'serve_utils.UpdateMode',
     pool: bool = False,
@@ -117,35 +123,63 @@ def apply(
     # Avoid circular import.
     from sky.client import sdk  # pylint: disable=import-outside-toplevel
 
-    dag = dag_utils.convert_entrypoint_to_dag(task)
-    with admin_policy_utils.apply_and_use_config_in_current_request(
-            dag, at_client_side=True) as dag:
-        sdk.validate(dag)
-        request_id = sdk.optimize(dag)
-        sdk.stream_and_get(request_id)
-        if _need_confirmation:
-            noun = 'pool' if pool else 'service'
-            prompt = f'Applying config to {noun} {service_name!r}. Proceed?'
-            if prompt is not None:
-                click.confirm(prompt,
-                              default=True,
-                              abort=True,
-                              show_default=True)
-
-        dag = client_common.upload_mounts_to_api_server(dag)
-        dag_str = dag_utils.dump_chain_dag_to_yaml_str(dag)
+    noun = 'pool' if pool else 'service'
+    # There are two cases here. If task is None, we should be trying to
+    # update the number of workers in the pool. If task is not None, we should
+    # be trying to apply a new config to the pool. The two code paths
+    # are slightly different with us needing to craft the dag and validate
+    # it if we have a task. In the future we could move this logic to the
+    # server side and simplify this code, for the time being we keep it here.
+    if task is None:
+        if workers is None:
+            raise ValueError(f'Cannot create a new {noun} without specifying '
+                             f'task or workers. Please provide either a task '
+                             f'or specify the number of workers.')
 
         body = payloads.JobsPoolApplyBody(
-            task=dag_str,
+            workers=workers,
             pool_name=service_name,
             mode=mode,
         )
+
         response = server_common.make_authenticated_request(
             'POST',
             '/jobs/pool_apply',
             json=json.loads(body.model_dump_json()),
             timeout=(5, None))
         return server_common.get_request_id(response)
+    else:
+        dag = dag_utils.convert_entrypoint_to_dag(task)
+        with admin_policy_utils.apply_and_use_config_in_current_request(
+                dag,
+                request_name=request_names.AdminPolicyRequestName.
+                JOBS_POOL_APPLY,
+                at_client_side=True) as dag:
+            sdk.validate(dag)
+            request_id = sdk.optimize(dag)
+            sdk.stream_and_get(request_id)
+            if _need_confirmation:
+                prompt = f'Applying config to {noun} {service_name!r}. Proceed?'
+                if prompt is not None:
+                    click.confirm(prompt,
+                                  default=True,
+                                  abort=True,
+                                  show_default=True)
+
+            dag = client_common.upload_mounts_to_api_server(dag)
+            dag_str = dag_utils.dump_chain_dag_to_yaml_str(dag)
+
+            body = payloads.JobsPoolApplyBody(
+                task=dag_str,
+                pool_name=service_name,
+                mode=mode,
+            )
+            response = server_common.make_authenticated_request(
+                'POST',
+                '/jobs/pool_apply',
+                json=json.loads(body.model_dump_json()),
+                timeout=(5, None))
+            return server_common.get_request_id(response)
 
 
 def down(

sky/serve/constants.py

@@ -65,7 +65,8 @@ AUTOSCALER_DEFAULT_DOWNSCALE_DELAY_SECONDS = 1200
 # TODO(tian): We might need to be careful that service logs can take a lot of
 # disk space. Maybe we could use a larger disk size, migrate to cloud storage or
 # do some log rotation.
-CONTROLLER_RESOURCES = {'cpus': '4+', 'disk_size': 200}
+# Set default minimal memory to 8GB to allow at least one service to run.
+CONTROLLER_RESOURCES = {'cpus': '4+', 'memory': '8+', 'disk_size': 200}
 # Autostop config for the jobs controller. These are the default values for
 # serve.controller.autostop in ~/.sky/config.yaml.
 CONTROLLER_AUTOSTOP = {
@@ -97,7 +98,7 @@ REPLICA_ID_ENV_VAR = 'SKYPILOT_SERVE_REPLICA_ID'
 # Changelog:
 # v1.0 - Introduce rolling update.
 # v2.0 - Added template-replica feature.
-# v3.0 - Added cluster pool.
+# v3.0 - Added pool.
 # v4.0 - Added pool argument to wait_service_registration.
 # v5.0 - Added pool argument to stream_serve_process_logs & stream_replica_logs.
 SERVE_VERSION = 5
@@ -106,7 +107,7 @@ TERMINATE_REPLICA_VERSION_MISMATCH_ERROR = (
     'The version of service is outdated and does not support manually '
     'terminating replicas. Please terminate the service and spin up again.')
 
-# Dummy run command for cluster pool.
+# Dummy run command for pool.
 POOL_DUMMY_RUN_COMMAND = 'echo "setup done"'
 
 # Error message for max number of services reached.

sky/serve/controller.py

@@ -21,7 +21,9 @@ from sky.serve import autoscalers
 from sky.serve import replica_managers
 from sky.serve import serve_state
 from sky.serve import serve_utils
+from sky.skylet import constants
 from sky.utils import common_utils
+from sky.utils import context_utils
 from sky.utils import ux_utils
 
 logger = sky_logging.init_logger(__name__)
@@ -44,13 +46,12 @@ class SkyServeController:
     """
 
     def __init__(self, service_name: str, service_spec: serve.SkyServiceSpec,
-                 service_task_yaml: str, host: str, port: int) -> None:
+                 version: int, host: str, port: int) -> None:
         self._service_name = service_name
         self._replica_manager: replica_managers.ReplicaManager = (
-            replica_managers.SkyPilotReplicaManager(
-                service_name=service_name,
-                spec=service_spec,
-                service_task_yaml_path=service_task_yaml))
+            replica_managers.SkyPilotReplicaManager(service_name=service_name,
+                                                    spec=service_spec,
+                                                    version=version))
         self._autoscaler: autoscalers.Autoscaler = (
             autoscalers.Autoscaler.from_spec(service_name, service_spec))
         self._host = host
@@ -172,7 +173,11 @@ class SkyServeController:
                 # See sky/serve/core.py::update
                 latest_task_yaml = serve_utils.generate_task_yaml_file_name(
                     self._service_name, version)
-                service = serve.SkyServiceSpec.from_yaml(latest_task_yaml)
+                with open(latest_task_yaml, 'r', encoding='utf-8') as f:
+                    yaml_content = f.read()
+                service = serve.SkyServiceSpec.from_yaml_str(yaml_content)
+                serve_state.add_or_update_version(self._service_name, version,
+                                                  service, yaml_content)
                 logger.info(
                     f'Update to new version version {version}: {service}')
 
@@ -283,9 +288,10 @@ class SkyServeController:
 # TODO(tian): Probably we should support service that will stop the VM in
 # specific time period.
 def run_controller(service_name: str, service_spec: serve.SkyServiceSpec,
-                   service_task_yaml: str, controller_host: str,
-                   controller_port: int):
-    controller = SkyServeController(service_name, service_spec,
-                                    service_task_yaml, controller_host,
-                                    controller_port)
+                   version: int, controller_host: str, controller_port: int):
+    os.environ[constants.OVERRIDE_CONSOLIDATION_MODE] = 'true'
+    # Hijack sys.stdout/stderr to be context aware.
+    context_utils.hijack_sys_attrs()
+    controller = SkyServeController(service_name, service_spec, version,
+                                    controller_host, controller_port)
     controller.run()

sky/serve/load_balancing_policies.py

@@ -121,7 +121,7 @@ class LeastLoadPolicy(LoadBalancingPolicy, name='least_load', default=True):
             return
         with self.lock:
             self.ready_replicas = ready_replicas
-            for r in self.ready_replicas:
+            for r in list(self.load_map.keys()):
                 if r not in ready_replicas:
                     del self.load_map[r]
             for replica in ready_replicas: