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

sky/serve/client/sdk.py

@@ -0,0 +1,366 @@
+"""SDK for SkyServe."""
+import json
+import typing
+from typing import List, Optional, Union
+
+import click
+import requests
+
+from sky.client import common as client_common
+from sky.server import common as server_common
+from sky.server.requests import payloads
+from sky.usage import usage_lib
+from sky.utils import dag_utils
+
+if typing.TYPE_CHECKING:
+    import io
+
+    import sky
+    from sky.serve import serve_utils
+
+
+@usage_lib.entrypoint
+@server_common.check_server_healthy_or_start
+def up(
+    task: Union['sky.Task', 'sky.Dag'],
+    service_name: str,
+    # Internal only:
+    # pylint: disable=invalid-name
+    _need_confirmation: bool = False
+) -> server_common.RequestId:
+    """Spins up a service.
+
+    Please refer to the sky.cli.serve_up for the document.
+
+    Args:
+        task: sky.Task to serve up.
+        service_name: Name of the service.
+        _need_confirmation: (Internal only) Whether to show a confirmation
+            prompt before spinning up the service.
+
+    Returns:
+        The request ID of the up request.
+
+    Request Returns:
+        service_name (str): The name of the service.  Same if passed in as an
+            argument.
+        endpoint (str): The service endpoint.
+    """
+
+    # Avoid circular import.
+    from sky.client import sdk  # pylint: disable=import-outside-toplevel
+
+    dag = dag_utils.convert_entrypoint_to_dag(task)
+    sdk.validate(dag)
+    request_id = sdk.optimize(dag)
+    sdk.stream_and_get(request_id)
+    if _need_confirmation:
+        prompt = f'Launching a new service {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.ServeUpBody(
+        task=dag_str,
+        service_name=service_name,
+    )
+    response = requests.post(
+        f'{server_common.get_server_url()}/serve/up',
+        json=json.loads(body.model_dump_json()),
+        timeout=(5, None),
+    )
+    return server_common.get_request_id(response)
+
+
+@usage_lib.entrypoint
+@server_common.check_server_healthy_or_start
+def update(
+    task: Union['sky.Task', 'sky.Dag'],
+    service_name: str,
+    mode: 'serve_utils.UpdateMode',
+    # Internal only:
+    # pylint: disable=invalid-name
+    _need_confirmation: bool = False
+) -> server_common.RequestId:
+    """Updates an existing service.
+
+    Please refer to the sky.cli.serve_update for the document.
+
+    Args:
+        task: sky.Task to update.
+        service_name: Name of the service.
+        mode: Update mode, including:
+            - sky.serve.UpdateMode.ROLLING
+            - sky.serve.UpdateMode.BLUE_GREEN
+        _need_confirmation: (Internal only) Whether to show a confirmation
+            prompt before updating the service.
+
+    Returns:
+        The request ID of the update request.
+
+    Request Returns:
+        None
+    """
+    # Avoid circular import.
+    from sky.client import sdk  # pylint: disable=import-outside-toplevel
+
+    dag = dag_utils.convert_entrypoint_to_dag(task)
+    sdk.validate(dag)
+    request_id = sdk.optimize(dag)
+    sdk.stream_and_get(request_id)
+    if _need_confirmation:
+        click.confirm(f'Updating service {service_name!r}. Proceed?',
+                      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.ServeUpdateBody(
+        task=dag_str,
+        service_name=service_name,
+        mode=mode,
+    )
+
+    response = requests.post(
+        f'{server_common.get_server_url()}/serve/update',
+        json=json.loads(body.model_dump_json()),
+        timeout=(5, None),
+    )
+    return server_common.get_request_id(response)
+
+
+@usage_lib.entrypoint
+@server_common.check_server_healthy_or_start
+def down(
+    service_names: Optional[Union[str, List[str]]],
+    all: bool = False,  # pylint: disable=redefined-builtin
+    purge: bool = False
+) -> server_common.RequestId:
+    """Tears down a service.
+
+    Please refer to the sky.cli.serve_down for the docs.
+
+    Args:
+        service_names: Name of the service(s).
+        all: Whether to terminate all services.
+        purge: Whether to terminate services in a failed status. These services
+          may potentially lead to resource leaks.
+
+    Returns:
+        The request ID of the down request.
+
+    Request Returns:
+        None
+
+    Request Raises:
+        sky.exceptions.ClusterNotUpError: if the sky serve controller is not up.
+        ValueError: if the arguments are invalid.
+        RuntimeError: if failed to terminate the service.
+    """
+    body = payloads.ServeDownBody(
+        service_names=service_names,
+        all=all,
+        purge=purge,
+    )
+    response = requests.post(
+        f'{server_common.get_server_url()}/serve/down',
+        json=json.loads(body.model_dump_json()),
+        timeout=(5, None),
+    )
+    return server_common.get_request_id(response)
+
+
+@usage_lib.entrypoint
+@server_common.check_server_healthy_or_start
+def terminate_replica(service_name: str, replica_id: int,
+                      purge: bool) -> server_common.RequestId:
+    """Tears down a specific replica for the given service.
+
+    Args:
+        service_name: Name of the service.
+        replica_id: ID of replica to terminate.
+        purge: Whether to terminate replicas in a failed status. These replicas
+          may lead to resource leaks, so we require the user to explicitly
+          specify this flag to make sure they are aware of this potential
+          resource leak.
+
+    Returns:
+        The request ID of the terminate replica request.
+
+    Request Raises:
+        sky.exceptions.ClusterNotUpError: if the sky sere controller is not up.
+        RuntimeError: if failed to terminate the replica.
+    """
+    body = payloads.ServeTerminateReplicaBody(
+        service_name=service_name,
+        replica_id=replica_id,
+        purge=purge,
+    )
+    response = requests.post(
+        f'{server_common.get_server_url()}/serve/terminate-replica',
+        json=json.loads(body.model_dump_json()),
+        timeout=(5, None),
+    )
+    return server_common.get_request_id(response)
+
+
+@usage_lib.entrypoint
+@server_common.check_server_healthy_or_start
+def status(
+        service_names: Optional[Union[str,
+                                      List[str]]]) -> server_common.RequestId:
+    """Gets service statuses.
+
+    If service_names is given, return those services. Otherwise, return all
+    services.
+
+    Each returned value has the following fields:
+
+    .. code-block:: python
+
+        {
+            'name': (str) service name,
+            'active_versions': (List[int]) a list of versions that are active,
+            'controller_job_id': (int) the job id of the controller,
+            'uptime': (int) uptime in seconds,
+            'status': (sky.ServiceStatus) service status,
+            'controller_port': (Optional[int]) controller port,
+            'load_balancer_port': (Optional[int]) load balancer port,
+            'endpoint': (Optional[str]) endpoint of the service,
+            'policy': (Optional[str]) autoscaling policy description,
+            'requested_resources_str': (str) str representation of
+              requested resources,
+            'load_balancing_policy': (str) load balancing policy name,
+            'replica_info': (List[Dict[str, Any]]) replica information,
+        }
+
+    Each entry in replica_info has the following fields:
+
+    .. code-block:: python
+
+        {
+            'replica_id': (int) replica id,
+            'name': (str) replica name,
+            'status': (sky.serve.ReplicaStatus) replica status,
+            'version': (int) replica version,
+            'launched_at': (int) timestamp of launched,
+            'handle': (ResourceHandle) handle of the replica cluster,
+            'endpoint': (str) endpoint of the replica,
+        }
+
+    For possible service statuses and replica statuses, please refer to
+    sky.cli.serve_status.
+
+    Args:
+        service_names: a single or a list of service names to query. If None,
+            query all services.
+
+    Returns:
+        The request ID of the status request.
+
+    Request Returns:
+        service_records (List[Dict[str, Any]]): A list of dicts, with each
+            dict containing the information of a service. If a service is not
+            found, it will be omitted from the returned list.
+
+    Request Raises:
+        RuntimeError: if failed to get the service status.
+        exceptions.ClusterNotUpError: if the sky serve controller is not up.
+    """
+    body = payloads.ServeStatusBody(service_names=service_names,)
+    response = requests.post(
+        f'{server_common.get_server_url()}/serve/status',
+        json=json.loads(body.model_dump_json()),
+        timeout=(5, None),
+    )
+    return server_common.get_request_id(response)
+
+
+@usage_lib.entrypoint
+@server_common.check_server_healthy_or_start
+def tail_logs(service_name: str,
+              target: Union[str, 'serve_utils.ServiceComponent'],
+              replica_id: Optional[int] = None,
+              follow: bool = True,
+              output_stream: Optional['io.TextIOBase'] = None) -> None:
+    """Tails logs for a service.
+
+    Usage:
+
+    .. code-block:: python
+
+        sky.serve.tail_logs(
+            service_name,
+            target=<component>,
+            follow=False, # Optionally, default to True
+            # replica_id=3, # Must be specified when target is REPLICA.
+        )
+
+
+    ``target`` is a enum of ``sky.serve.ServiceComponent``, which can be one of:
+
+    - ``sky.serve.ServiceComponent.CONTROLLER``
+
+    - ``sky.serve.ServiceComponent.LOAD_BALANCER``
+
+    - ``sky.serve.ServiceComponent.REPLICA``
+
+    Pass target as a lower-case string is also supported, e.g.
+    ``target='controller'``.
+    To use ``sky.serve.ServiceComponent.REPLICA``, you must specify
+    ``replica_id``.
+
+    To tail controller logs:
+
+    .. code-block:: python
+
+        # follow default to True
+        sky.serve.tail_logs(
+            service_name, target=sky.serve.ServiceComponent.CONTROLLER
+        )
+
+    To print replica 3 logs:
+
+    .. code-block:: python
+
+        # Pass target as a lower-case string is also supported.
+        sky.serve.tail_logs(
+            service_name, target='replica',
+            follow=False, replica_id=3
+        )
+
+    Args:
+        service_name: Name of the service.
+        target: The component to tail logs.
+        replica_id: The ID of the replica to tail logs.
+        follow: Whether to follow the logs.
+        output_stream: The stream to write the logs to. If None, print to the
+            console.
+
+    Returns:
+        The request ID of the tail logs request.
+
+    Request Raises:
+        sky.exceptions.ClusterNotUpError: the sky serve controller is not up.
+        ValueError: arguments not valid, or failed to tail the logs.
+    """
+    # Avoid circular import.
+    from sky.client import sdk  # pylint: disable=import-outside-toplevel
+
+    body = payloads.ServeLogsBody(
+        service_name=service_name,
+        target=target,
+        replica_id=replica_id,
+        follow=follow,
+    )
+    response = requests.post(
+        f'{server_common.get_server_url()}/serve/logs',
+        json=json.loads(body.model_dump_json()),
+        timeout=(5, None),
+        stream=True,
+    )
+    request_id = server_common.get_request_id(response)
+    sdk.stream_response(request_id, response, output_stream)

sky/serve/constants.py

@@ -12,6 +12,9 @@ PORT_SELECTION_FILE_LOCK_PATH = f'{SKYSERVE_METADATA_DIR}/port_selection.lock'
 # Signal file path for controller to handle signals.
 SIGNAL_FILE_PATH = '/tmp/sky_serve_controller_signal_{}'
 
+# Time to wait in seconds for controller to setup, this involves the time to run
+# cloud dependencies installation.
+CONTROLLER_SETUP_TIMEOUT_SECONDS = 300
 # Time to wait in seconds for service to register on the controller.
 SERVICE_REGISTER_TIMEOUT_SECONDS = 60
 
@@ -39,8 +42,7 @@ ENDPOINT_PROBE_INTERVAL_SECONDS = 10
 # The default timeout in seconds for a readiness probe request. We set the
 # timeout to 15s since using actual generation in LLM services as readiness
 # probe is very time-consuming (33B, 70B, ...).
-# TODO(tian): Expose this option to users in yaml file.
-READINESS_PROBE_TIMEOUT_SECONDS = 15
+DEFAULT_READINESS_PROBE_TIMEOUT_SECONDS = 15
 
 # Autoscaler window size in seconds for query per second. We calculate qps by
 # divide the number of queries in last window size by this window size.
@@ -93,4 +95,11 @@ REPLICA_ID_ENV_VAR = 'SKYPILOT_SERVE_REPLICA_ID'
 # change for the serve_utils.ServeCodeGen, we need to bump this version, so that
 # the user can be notified to update their SkyPilot serve version on the remote
 # cluster.
-SERVE_VERSION = 1
+# Changelog:
+# v1.0 - Introduce rolling update.
+# v2.0 - Added template-replica feature.
+SERVE_VERSION = 2
+
+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.')

sky/serve/controller.py

@@ -2,14 +2,16 @@
 
 Responsible for autoscaling and replica management.
 """
+import contextlib
 import logging
-import os
 import threading
 import time
 import traceback
 from typing import Any, Dict, List
 
+import colorama
 import fastapi
+from fastapi import responses
 import uvicorn
 
 from sky import serve
@@ -50,7 +52,14 @@ class SkyServeController:
             autoscalers.Autoscaler.from_spec(service_name, service_spec))
         self._host = host
         self._port = port
-        self._app = fastapi.FastAPI()
+        self._app = fastapi.FastAPI(lifespan=self.lifespan)
+
+    @contextlib.asynccontextmanager
+    async def lifespan(self, _: fastapi.FastAPI):
+        uvicorn_access_logger = logging.getLogger('uvicorn.access')
+        for handler in uvicorn_access_logger.handlers:
+            handler.setFormatter(sky_logging.FORMATTER)
+        yield
 
     def _run_autoscaler(self):
         logger.info('Starting autoscaler.')
@@ -58,9 +67,16 @@ class SkyServeController:
             try:
                 replica_infos = serve_state.get_replica_infos(
                     self._service_name)
+                # Use the active versions set by replica manager to make
+                # sure we only scale down the outdated replicas that are
+                # not used by the load balancer.
+                record = serve_state.get_service_from_name(self._service_name)
+                assert record is not None, ('No service record found for '
+                                            f'{self._service_name}')
+                active_versions = record['active_versions']
                 logger.info(f'All replica info: {replica_infos}')
-                scaling_options = self._autoscaler.evaluate_scaling(
-                    replica_infos)
+                scaling_options = self._autoscaler.generate_scaling_decisions(
+                    replica_infos, active_versions)
                 for scaling_option in scaling_options:
                     logger.info(f'Scaling option received: {scaling_option}')
                     if (scaling_option.operator ==
@@ -68,15 +84,10 @@ class SkyServeController:
                         assert (scaling_option.target is None or isinstance(
                             scaling_option.target, dict)), scaling_option
                         self._replica_manager.scale_up(scaling_option.target)
-                    elif (scaling_option.operator ==
-                          autoscalers.AutoscalerDecisionOperator.SCALE_DOWN):
+                    else:
                         assert isinstance(scaling_option.target,
                                           int), scaling_option
                         self._replica_manager.scale_down(scaling_option.target)
-                    else:
-                        with ux_utils.enable_traceback():
-                            logger.error('Error in scaling_option.operator: '
-                                         f'{scaling_option.operator}')
             except Exception as e:  # pylint: disable=broad-except
                 # No matter what error happens, we should keep the
                 # monitor running.
@@ -89,7 +100,8 @@ class SkyServeController:
     def run(self) -> None:
 
         @self._app.post('/controller/load_balancer_sync')
-        async def load_balancer_sync(request: fastapi.Request):
+        async def load_balancer_sync(
+                request: fastapi.Request) -> fastapi.Response:
             request_data = await request.json()
             # TODO(MaoZiming): Check aggregator type.
             request_aggregator: Dict[str, Any] = request_data.get(
@@ -97,18 +109,21 @@ class SkyServeController:
             timestamps: List[int] = request_aggregator.get('timestamps', [])
             logger.info(f'Received {len(timestamps)} inflight requests.')
             self._autoscaler.collect_request_information(request_aggregator)
-            return {
+            return responses.JSONResponse(content={
                 'ready_replica_urls':
                     self._replica_manager.get_active_replica_urls()
-            }
+            },
+                                          status_code=200)
 
         @self._app.post('/controller/update_service')
-        async def update_service(request: fastapi.Request):
+        async def update_service(request: fastapi.Request) -> fastapi.Response:
             request_data = await request.json()
             try:
                 version = request_data.get('version', None)
                 if version is None:
-                    return {'message': 'Error: version is not specified.'}
+                    return responses.JSONResponse(
+                        content={'message': 'Error: version is not specified.'},
+                        status_code=400)
                 update_mode_str = request_data.get(
                     'mode', serve_utils.DEFAULT_UPDATE_MODE.value)
                 update_mode = serve_utils.UpdateMode(update_mode_str)
@@ -137,40 +152,95 @@ class SkyServeController:
                 self._autoscaler.update_version(version,
                                                 service,
                                                 update_mode=update_mode)
-                return {'message': 'Success'}
+                return responses.JSONResponse(content={'message': 'Success'},
+                                              status_code=200)
             except Exception as e:  # pylint: disable=broad-except
                 logger.error(f'Error in update_service: '
                              f'{common_utils.format_exception(e)}')
-                return {'message': 'Error'}
+                return responses.JSONResponse(content={'message': 'Error'},
+                                              status_code=500)
 
-        @self._app.on_event('startup')
-        def configure_logger():
-            uvicorn_access_logger = logging.getLogger('uvicorn.access')
-            for handler in uvicorn_access_logger.handlers:
-                handler.setFormatter(sky_logging.FORMATTER)
+        @self._app.post('/controller/terminate_replica')
+        async def terminate_replica(
+                request: fastapi.Request) -> fastapi.Response:
+            request_data = await request.json()
+            replica_id = request_data['replica_id']
+            assert isinstance(replica_id,
+                              int), 'Error: replica ID must be an integer.'
+            purge = request_data['purge']
+            assert isinstance(purge, bool), 'Error: purge must be a boolean.'
+            replica_info = serve_state.get_replica_info_from_id(
+                self._service_name, replica_id)
+            assert replica_info is not None, (f'Error: replica '
+                                              f'{replica_id} does not exist.')
+            replica_status = replica_info.status
+
+            if replica_status == serve_state.ReplicaStatus.SHUTTING_DOWN:
+                return responses.JSONResponse(
+                    status_code=409,
+                    content={
+                        'message':
+                            f'Replica {replica_id} of service '
+                            f'{self._service_name!r} is already in the process '
+                            f'of terminating. Skip terminating now.'
+                    })
+
+            if (replica_status in serve_state.ReplicaStatus.failed_statuses()
+                    and not purge):
+                return responses.JSONResponse(
+                    status_code=409,
+                    content={
+                        'message': f'{colorama.Fore.YELLOW}Replica '
+                                   f'{replica_id} of service '
+                                   f'{self._service_name!r} is in failed '
+                                   f'status ({replica_info.status}). '
+                                   f'Skipping its termination as it could '
+                                   f'lead to a resource leak. '
+                                   f'(Use `sky serve down '
+                                   f'{self._service_name!r} --replica-id '
+                                   f'{replica_id} --purge` to '
+                                   'forcefully terminate the replica.)'
+                                   f'{colorama.Style.RESET_ALL}'
+                    })
+
+            self._replica_manager.scale_down(replica_id, purge=purge)
+
+            action = 'terminated' if not purge else 'purged'
+            message = (f'{colorama.Fore.GREEN}Replica {replica_id} of service '
+                       f'{self._service_name!r} is scheduled to be '
+                       f'{action}.{colorama.Style.RESET_ALL}\n'
+                       f'Please use {ux_utils.BOLD}sky serve status '
+                       f'{self._service_name}{ux_utils.RESET_BOLD} '
+                       f'to check the latest status.')
+            return responses.JSONResponse(status_code=200,
+                                          content={'message': message})
+
+        @self._app.exception_handler(Exception)
+        async def validation_exception_handler(
+                request: fastapi.Request, exc: Exception) -> fastapi.Response:
+            with ux_utils.enable_traceback():
+                logger.error(f'Error in controller: {exc!r}')
+            return responses.JSONResponse(
+                status_code=500,
+                content={
+                    'message':
+                        (f'Failed method {request.method} at URL {request.url}.'
+                         f' Exception message is {exc!r}.')
+                },
+            )
 
         threading.Thread(target=self._run_autoscaler).start()
 
         logger.info('SkyServe Controller started on '
                     f'http://{self._host}:{self._port}')
 
-        uvicorn.run(self._app, host={self._host}, port=self._port)
+        uvicorn.run(self._app, host=self._host, port=self._port)
 
 
 # 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,
-                   task_yaml: str, controller_port: int):
-    # We expose the controller to the public network when running inside a
-    # kubernetes cluster to allow external load balancers (example, for
-    # high availability load balancers) to communicate with the controller.
-    def _get_host():
-        if 'KUBERNETES_SERVICE_HOST' in os.environ:
-            return '0.0.0.0'
-        else:
-            return 'localhost'
-
-    host = _get_host()
-    controller = SkyServeController(service_name, service_spec, task_yaml, host,
-                                    controller_port)
+                   task_yaml: str, controller_host: str, controller_port: int):
+    controller = SkyServeController(service_name, service_spec, task_yaml,
+                                    controller_host, controller_port)
     controller.run()