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

sky/serve/server/impl.py

@@ -0,0 +1,1093 @@
+"""Implementation of the SkyServe core APIs."""
+import pathlib
+import re
+import shlex
+import signal
+import tempfile
+import threading
+import typing
+from typing import Any, Dict, List, Optional, Set, Tuple, Union
+import uuid
+
+import colorama
+import filelock
+
+from sky import backends
+from sky import exceptions
+from sky import execution
+from sky import sky_logging
+from sky import skypilot_config
+from sky import task as task_lib
+from sky.adaptors import common as adaptors_common
+from sky.backends import backend_utils
+from sky.catalog import common as service_catalog_common
+from sky.data import storage as storage_lib
+from sky.serve import constants as serve_constants
+from sky.serve import serve_rpc_utils
+from sky.serve import serve_state
+from sky.serve import serve_utils
+from sky.server.requests import request_names
+from sky.skylet import constants
+from sky.skylet import job_lib
+from sky.utils import admin_policy_utils
+from sky.utils import command_runner
+from sky.utils import common
+from sky.utils import common_utils
+from sky.utils import controller_utils
+from sky.utils import dag_utils
+from sky.utils import rich_utils
+from sky.utils import subprocess_utils
+from sky.utils import ux_utils
+from sky.utils import yaml_utils
+
+if typing.TYPE_CHECKING:
+    import grpc
+else:
+    grpc = adaptors_common.LazyImport('grpc')
+
+logger = sky_logging.init_logger(__name__)
+
+
+def _rewrite_tls_credential_paths_and_get_tls_env_vars(
+        service_name: str, task: 'task_lib.Task') -> Dict[str, Any]:
+    """Rewrite the paths of TLS credentials in the task.
+
+    Args:
+        service_name: Name of the service.
+        task: sky.Task to rewrite.
+
+    Returns:
+        The generated template variables for TLS.
+    """
+    service_spec = task.service
+    # Already checked by validate_service_task
+    assert service_spec is not None
+    if service_spec.tls_credential is None:
+        return {'use_tls': False}
+    remote_tls_keyfile = (
+        serve_utils.generate_remote_tls_keyfile_name(service_name))
+    remote_tls_certfile = (
+        serve_utils.generate_remote_tls_certfile_name(service_name))
+    tls_template_vars = {
+        'use_tls': True,
+        'remote_tls_keyfile': remote_tls_keyfile,
+        'remote_tls_certfile': remote_tls_certfile,
+        'local_tls_keyfile': service_spec.tls_credential.keyfile,
+        'local_tls_certfile': service_spec.tls_credential.certfile,
+    }
+    service_spec.tls_credential = serve_utils.TLSCredential(
+        remote_tls_keyfile, remote_tls_certfile)
+    return tls_template_vars
+
+
+def _get_service_record(
+        service_name: str, pool: bool,
+        handle: backends.CloudVmRayResourceHandle,
+        backend: backends.CloudVmRayBackend) -> Optional[Dict[str, Any]]:
+    """Get the service record."""
+    noun = 'pool' if pool else 'service'
+
+    assert isinstance(handle, backends.CloudVmRayResourceHandle)
+    use_legacy = not handle.is_grpc_enabled_with_flag
+
+    if not use_legacy:
+        try:
+            service_statuses = serve_rpc_utils.RpcRunner.get_service_status(
+                handle, [service_name], pool)
+        except exceptions.SkyletMethodNotImplementedError:
+            use_legacy = True
+
+    if use_legacy:
+        code = serve_utils.ServeCodeGen.get_service_status([service_name],
+                                                           pool=pool)
+        returncode, serve_status_payload, stderr = backend.run_on_head(
+            handle,
+            code,
+            require_outputs=True,
+            stream_logs=False,
+            separate_stderr=True)
+        try:
+            subprocess_utils.handle_returncode(returncode,
+                                               code,
+                                               f'Failed to get {noun} status',
+                                               stderr,
+                                               stream_logs=True)
+        except exceptions.CommandError as e:
+            raise RuntimeError(e.error_msg) from e
+
+        service_statuses = serve_utils.load_service_status(serve_status_payload)
+
+    assert len(service_statuses) <= 1, service_statuses
+    if not service_statuses:
+        return None
+    return service_statuses[0]
+
+
+def _maybe_display_run_warning(task: 'task_lib.Task') -> None:
+    # We do not block the user from creating a pool with a run section
+    # in order to enable using the same yaml for pool creation
+    # and job submission. But we want to make it clear that 'run' will not
+    # be respected here.
+    if task.run is not None:
+        logger.warning(
+            f'{colorama.Fore.YELLOW} Pool creation does not support the '
+            '`run` section. Creating the pool while ignoring the '
+            f'`run` section.{colorama.Style.RESET_ALL}')
+
+
+def up(
+    task: 'task_lib.Task',
+    service_name: Optional[str] = None,
+    pool: bool = False,
+) -> Tuple[str, str]:
+    """Spins up a service or a pool."""
+    task.validate()
+    serve_utils.validate_service_task(task, pool=pool)
+    assert task.service is not None
+    assert task.service.pool == pool, 'Inconsistent pool flag.'
+    noun = 'pool' if pool else 'service'
+    capnoun = noun.capitalize()
+    if service_name is None:
+        service_name = serve_utils.generate_service_name(pool)
+
+    # The service name will be used as:
+    # 1. controller cluster name: 'sky-serve-controller-<service_name>'
+    # 2. replica cluster name: '<service_name>-<replica_id>'
+    # In both cases, service name shares the same regex with cluster name.
+    if re.fullmatch(constants.CLUSTER_NAME_VALID_REGEX, service_name) is None:
+        with ux_utils.print_exception_no_traceback():
+            raise ValueError(f'{capnoun} name {service_name!r} is invalid: '
+                             f'ensure it is fully matched by regex (e.g., '
+                             'only contains lower letters, numbers and dash): '
+                             f'{constants.CLUSTER_NAME_VALID_REGEX}')
+
+    dag = dag_utils.convert_entrypoint_to_dag(task)
+    # Always apply the policy again here, even though it might have been applied
+    # in the CLI. This is to ensure that we apply the policy to the final DAG
+    # and get the mutated config.
+    dag, mutated_user_config = admin_policy_utils.apply(
+        dag, request_name=request_names.AdminPolicyRequestName.SERVE_UP)
+    dag.resolve_and_validate_volumes()
+    dag.pre_mount_volumes()
+    task = dag.tasks[0]
+    assert task.service is not None
+    if pool:
+        _maybe_display_run_warning(task)
+        # Use dummy run script for pool.
+        task.run = serve_constants.POOL_DUMMY_RUN_COMMAND
+
+    with rich_utils.safe_status(
+            ux_utils.spinner_message(f'Initializing {noun}')):
+        # Handle file mounts using two-hop approach when cloud storage
+        # unavailable
+        storage_clouds = (
+            storage_lib.get_cached_enabled_storage_cloud_names_or_refresh())
+        force_disable_cloud_bucket = skypilot_config.get_nested(
+            ('serve', 'force_disable_cloud_bucket'), False)
+        if storage_clouds and not force_disable_cloud_bucket:
+            controller_utils.maybe_translate_local_file_mounts_and_sync_up(
+                task, task_type='serve')
+            local_to_controller_file_mounts = {}
+        else:
+            # Fall back to two-hop file_mount uploading when no cloud storage
+            if task.storage_mounts:
+                raise exceptions.NotSupportedError(
+                    'Cloud-based file_mounts are specified, but no cloud '
+                    'storage is available. Please specify local '
+                    'file_mounts only.')
+            local_to_controller_file_mounts = (
+                controller_utils.translate_local_file_mounts_to_two_hop(task))
+
+    tls_template_vars = _rewrite_tls_credential_paths_and_get_tls_env_vars(
+        service_name, task)
+
+    with tempfile.NamedTemporaryFile(
+            prefix=f'service-task-{service_name}-',
+            mode='w',
+    ) as service_file, tempfile.NamedTemporaryFile(
+            prefix=f'controller-task-{service_name}-',
+            mode='w',
+    ) as controller_file:
+        controller = controller_utils.get_controller_for_pool(pool)
+        controller_name = controller.value.cluster_name
+        task_config = task.to_yaml_config()
+        yaml_utils.dump_yaml(service_file.name, task_config)
+        remote_tmp_task_yaml_path = (
+            serve_utils.generate_remote_tmp_task_yaml_file_name(service_name))
+        remote_config_yaml_path = (
+            serve_utils.generate_remote_config_yaml_file_name(service_name))
+        controller_log_file = (
+            serve_utils.generate_remote_controller_log_file_name(service_name))
+        controller_resources = controller_utils.get_controller_resources(
+            controller=controller, task_resources=task.resources)
+        controller_job_id = None
+        if serve_utils.is_consolidation_mode(pool):
+            # We need a unique integer per sky.serve.up call to avoid name
+            # conflict. Originally in non-consolidation mode, this is the ray
+            # job id; now we use the request id hash instead. Here we also
+            # make sure it is a 32-bit integer to avoid overflow on sqlalchemy.
+            rid = common_utils.get_current_request_id()
+            controller_job_id = hash(uuid.UUID(rid).int) & 0x7FFFFFFF
+
+        vars_to_fill = {
+            'remote_task_yaml_path': remote_tmp_task_yaml_path,
+            'local_task_yaml_path': service_file.name,
+            'service_name': service_name,
+            'controller_log_file': controller_log_file,
+            'remote_user_config_path': remote_config_yaml_path,
+            'local_to_controller_file_mounts': local_to_controller_file_mounts,
+            'modified_catalogs':
+                service_catalog_common.get_modified_catalog_file_mounts(),
+            'consolidation_mode_job_id': controller_job_id,
+            'entrypoint': shlex.quote(common_utils.get_current_command()),
+            **tls_template_vars,
+            **controller_utils.shared_controller_vars_to_fill(
+                controller=controller_utils.Controllers.SKY_SERVE_CONTROLLER,
+                remote_user_config_path=remote_config_yaml_path,
+                local_user_config=mutated_user_config,
+            ),
+        }
+        common_utils.fill_template(serve_constants.CONTROLLER_TEMPLATE,
+                                   vars_to_fill,
+                                   output_path=controller_file.name)
+        controller_task = task_lib.Task.from_yaml(controller_file.name)
+        # TODO(tian): Probably run another sky.launch after we get the load
+        # balancer port from the controller? So we don't need to open so many
+        # ports here. Or, we should have a nginx traffic control to refuse
+        # any connection to the unregistered ports.
+        if not pool:
+            controller_resources = {
+                r.copy(ports=[serve_constants.LOAD_BALANCER_PORT_RANGE])
+                for r in controller_resources
+            }
+        controller_task.set_resources(controller_resources)
+
+        # # Set service_name so the backend will know to modify default ray
+        # task CPU usage to custom value instead of default 0.5 vCPU. We need
+        # to set it to a smaller value to support a larger number of services.
+        controller_task.service_name = service_name
+
+        # We directly submit the request to the controller and let the
+        # controller to check name conflict. Suppose we have multiple
+        # sky.serve.up() with same service name, the first one will
+        # successfully write its job id to controller service database;
+        # and for all following sky.serve.up(), the controller will throw
+        # an exception (name conflict detected) and exit. Therefore the
+        # controller job id in database could be use as an indicator of
+        # whether the service is already running. If the id is the same
+        # with the current job id, we know the service is up and running
+        # for the first time; otherwise it is a name conflict.
+        # Since the controller may be shared among multiple users, launch the
+        # controller with the API server's user hash.
+        if not serve_utils.is_consolidation_mode(pool):
+            print(f'{colorama.Fore.YELLOW}Launching controller for '
+                  f'{service_name!r}...{colorama.Style.RESET_ALL}')
+            with common.with_server_user():
+                with skypilot_config.local_active_workspace_ctx(
+                        constants.SKYPILOT_DEFAULT_WORKSPACE):
+                    controller_job_id, controller_handle = execution.launch(
+                        task=controller_task,
+                        cluster_name=controller_name,
+                        retry_until_up=True,
+                        _request_name=request_names.AdminPolicyRequestName.
+                        SERVE_LAUNCH_CONTROLLER,
+                        _disable_controller_check=True,
+                    )
+        else:
+            controller_type = controller_utils.get_controller_for_pool(pool)
+            controller_handle = backend_utils.is_controller_accessible(
+                controller=controller_type, stopped_message='')
+            backend = backend_utils.get_backend_from_handle(controller_handle)
+            assert isinstance(backend, backends.CloudVmRayBackend)
+            backend.sync_file_mounts(
+                handle=controller_handle,
+                all_file_mounts=controller_task.file_mounts,
+                storage_mounts=controller_task.storage_mounts)
+            run_script = controller_task.run
+            assert isinstance(run_script, str)
+            # Manually add the env variables to the run script. Originally
+            # this is done in ray jobs submission but now we have to do it
+            # manually because there is no ray runtime on the API server.
+            env_cmds = [
+                f'export {k}={v!r}' for k, v in controller_task.envs.items()
+            ]
+            run_script = '\n'.join(env_cmds + [run_script])
+            # Dump script for high availability recovery.
+            serve_state.set_ha_recovery_script(service_name, run_script)
+            backend.run_on_head(controller_handle, run_script)
+
+        style = colorama.Style
+        fore = colorama.Fore
+
+        assert controller_job_id is not None and controller_handle is not None
+        assert isinstance(controller_handle, backends.CloudVmRayResourceHandle)
+        backend = backend_utils.get_backend_from_handle(controller_handle)
+        assert isinstance(backend, backends.CloudVmRayBackend)
+        # TODO(tian): Cache endpoint locally to speedup. Endpoint won't
+        # change after the first time, so there is no consistency issue.
+        try:
+            with rich_utils.safe_status(
+                    ux_utils.spinner_message(
+                        f'Waiting for the {noun} to register')):
+                # This function will check the controller job id in the database
+                # and return the endpoint if the job id matches. Otherwise it
+                # will return None.
+                use_legacy = not controller_handle.is_grpc_enabled_with_flag
+
+                if controller_handle.is_grpc_enabled_with_flag:
+                    try:
+                        lb_port = serve_rpc_utils.RpcRunner.wait_service_registration(  # pylint: disable=line-too-long
+                            controller_handle, service_name, controller_job_id,
+                            pool)
+                    except exceptions.SkyletMethodNotImplementedError:
+                        use_legacy = True
+
+                if use_legacy:
+                    code = serve_utils.ServeCodeGen.wait_service_registration(
+                        service_name, controller_job_id, pool)
+                    returncode, lb_port_payload, _ = backend.run_on_head(
+                        controller_handle,
+                        code,
+                        require_outputs=True,
+                        stream_logs=False)
+                    subprocess_utils.handle_returncode(
+                        returncode, code,
+                        f'Failed to wait for {noun} initialization',
+                        lb_port_payload)
+                    lb_port = serve_utils.load_service_initialization_result(
+                        lb_port_payload)
+        except (exceptions.CommandError, grpc.FutureTimeoutError,
+                grpc.RpcError):
+            if serve_utils.is_consolidation_mode(pool):
+                with ux_utils.print_exception_no_traceback():
+                    raise RuntimeError(
+                        f'Failed to wait for {noun} initialization. '
+                        'Please check the logs above for more details.'
+                    ) from None
+            statuses = backend.get_job_status(controller_handle,
+                                              [controller_job_id],
+                                              stream_logs=False)
+            controller_job_status = list(statuses.values())[0]
+            if controller_job_status == job_lib.JobStatus.PENDING:
+                # Max number of services reached due to vCPU constraint.
+                # The controller job is pending due to ray job scheduling.
+                # We manually cancel the job here.
+                backend.cancel_jobs(controller_handle, [controller_job_id])
+                with ux_utils.print_exception_no_traceback():
+                    raise RuntimeError(
+                        'Max number of services reached. '
+                        'To spin up more services, please '
+                        'tear down some existing services.') from None
+            else:
+                # Possible cases:
+                # (1) name conflict;
+                # (2) max number of services reached due to memory
+                # constraint. The job will successfully run on the
+                # controller, but there will be an error thrown due
+                # to memory constraint check in the controller.
+                # See sky/serve/service.py for more details.
+                with ux_utils.print_exception_no_traceback():
+                    raise RuntimeError(
+                        'Failed to spin up the service. Please '
+                        'check the logs above for more details.') from None
+        else:
+            if not serve_utils.is_consolidation_mode(pool) and not pool:
+                socket_endpoint = backend_utils.get_endpoints(
+                    controller_handle.cluster_name,
+                    lb_port,
+                    skip_status_check=True).get(lb_port)
+            else:
+                socket_endpoint = f'localhost:{lb_port}'
+            assert socket_endpoint is not None, (
+                'Did not get endpoint for controller.')
+            # Already checked by validate_service_task
+            assert task.service is not None
+            protocol = ('http'
+                        if task.service.tls_credential is None else 'https')
+            socket_endpoint = socket_endpoint.replace('https://', '').replace(
+                'http://', '')
+            endpoint = f'{protocol}://{socket_endpoint}'
+
+        if pool:
+            logger.info(
+                f'{fore.CYAN}Pool name: '
+                f'{style.BRIGHT}{service_name}{style.RESET_ALL}'
+                f'\n📋 Useful Commands'
+                f'\n{ux_utils.INDENT_SYMBOL}To submit jobs to the pool:\t'
+                f'{ux_utils.BOLD}sky jobs launch --pool {service_name} '
+                f'<yaml_file>{ux_utils.RESET_BOLD}'
+                f'\n{ux_utils.INDENT_SYMBOL}To submit multiple jobs:\t'
+                f'{ux_utils.BOLD}sky jobs launch --pool {service_name} '
+                f'--num-jobs 10 <yaml_file>{ux_utils.RESET_BOLD}'
+                f'\n{ux_utils.INDENT_SYMBOL}To check the pool status:\t'
+                f'{ux_utils.BOLD}sky jobs pool status {service_name}'
+                f'{ux_utils.RESET_BOLD}'
+                f'\n{ux_utils.INDENT_LAST_SYMBOL}To terminate the pool:\t'
+                f'{ux_utils.BOLD}sky jobs pool down {service_name}'
+                f'{ux_utils.RESET_BOLD}'
+                f'\n{ux_utils.INDENT_SYMBOL}To update the number of workers:\t'
+                f'{ux_utils.BOLD}sky jobs pool apply --pool {service_name} '
+                f'--workers 5{ux_utils.RESET_BOLD}'
+                '\n\n' + ux_utils.finishing_message('Successfully created pool '
+                                                    f'{service_name!r}.'))
+        else:
+            logger.info(
+                f'{fore.CYAN}Service name: '
+                f'{style.BRIGHT}{service_name}{style.RESET_ALL}'
+                f'\n{fore.CYAN}Endpoint URL: '
+                f'{style.BRIGHT}{endpoint}{style.RESET_ALL}'
+                f'\n📋 Useful Commands'
+                f'\n{ux_utils.INDENT_SYMBOL}To check service status:\t'
+                f'{ux_utils.BOLD}sky serve status {service_name} '
+                f'[--endpoint]{ux_utils.RESET_BOLD}'
+                f'\n{ux_utils.INDENT_SYMBOL}To teardown the service:\t'
+                f'{ux_utils.BOLD}sky serve down {service_name}'
+                f'{ux_utils.RESET_BOLD}'
+                f'\n{ux_utils.INDENT_SYMBOL}To see replica logs:\t'
+                f'{ux_utils.BOLD}sky serve logs {service_name} [REPLICA_ID]'
+                f'{ux_utils.RESET_BOLD}'
+                f'\n{ux_utils.INDENT_SYMBOL}To see load balancer logs:\t'
+                f'{ux_utils.BOLD}sky serve logs --load-balancer {service_name}'
+                f'{ux_utils.RESET_BOLD}'
+                f'\n{ux_utils.INDENT_SYMBOL}To see controller logs:\t'
+                f'{ux_utils.BOLD}sky serve logs --controller {service_name}'
+                f'{ux_utils.RESET_BOLD}'
+                f'\n{ux_utils.INDENT_SYMBOL}To monitor the status:\t'
+                f'{ux_utils.BOLD}watch -n10 sky serve status {service_name}'
+                f'{ux_utils.RESET_BOLD}'
+                f'\n{ux_utils.INDENT_LAST_SYMBOL}To send a test request:\t'
+                f'{ux_utils.BOLD}curl {endpoint}'
+                f'{ux_utils.RESET_BOLD}'
+                '\n\n' + ux_utils.finishing_message(
+                    'Service is spinning up and replicas '
+                    'will be ready shortly.'))
+        return service_name, endpoint
+
+
+def update(
+    task: Optional['task_lib.Task'],
+    service_name: str,
+    mode: serve_utils.UpdateMode = serve_utils.DEFAULT_UPDATE_MODE,
+    pool: bool = False,
+    workers: Optional[int] = None,
+) -> None:
+    """Updates an existing service or pool."""
+    noun = 'pool' if pool else 'service'
+    capnoun = noun.capitalize()
+
+    controller_type = controller_utils.get_controller_for_pool(pool)
+    handle = backend_utils.is_controller_accessible(
+        controller=controller_type,
+        stopped_message=
+        'Service controller is stopped. There is no service to update. '
+        f'To spin up a new service, use {ux_utils.BOLD}'
+        f'sky serve up{ux_utils.RESET_BOLD}',
+        non_existent_message='Service does not exist. '
+        'To spin up a new service, '
+        f'use {ux_utils.BOLD}sky serve up{ux_utils.RESET_BOLD}',
+    )
+
+    assert isinstance(handle, backends.CloudVmRayResourceHandle)
+    backend = backend_utils.get_backend_from_handle(handle)
+    assert isinstance(backend, backends.CloudVmRayBackend)
+
+    service_record = _get_service_record(service_name, pool, handle, backend)
+
+    if service_record is None:
+        cmd = 'sky jobs pool up' if pool else 'sky serve up'
+        with ux_utils.print_exception_no_traceback():
+            raise RuntimeError(f'Cannot find {noun} {service_name!r}.'
+                               f'To spin up a {noun}, use {ux_utils.BOLD}'
+                               f'{cmd}{ux_utils.RESET_BOLD}')
+
+    # If task is None and workers is specified, load existing configuration
+    # and update replica count.
+    if task is None:
+        if workers is None:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(
+                    f'Cannot update {noun} without specifying '
+                    f'task or workers. Please provide either a task '
+                    f'or specify the number of workers.')
+
+        if not pool:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(
+                    'Non-pool service, trying to update replicas to '
+                    f'{workers} is not supported. Ignoring the update.')
+
+        # Load the existing task configuration from the service's YAML file
+        yaml_content = service_record['yaml_content']
+
+        # Load the existing task configuration
+        task = task_lib.Task.from_yaml_str(yaml_content)
+
+        if task.service is None:
+            with ux_utils.print_exception_no_traceback():
+                raise RuntimeError('No service configuration found in '
+                                   f'existing {noun} {service_name!r}')
+        task.set_service(task.service.copy(min_replicas=workers))
+
+    task.validate()
+    serve_utils.validate_service_task(task, pool=pool)
+
+    # Now apply the policy and handle task-specific logic
+    # Always apply the policy again here, even though it might have been applied
+    # in the CLI. This is to ensure that we apply the policy to the final DAG
+    # and get the mutated config.
+    # TODO(cblmemo,zhwu): If a user sets a new skypilot_config, the update
+    # will not apply the config.
+    dag, _ = admin_policy_utils.apply(
+        task, request_name=request_names.AdminPolicyRequestName.SERVE_UPDATE)
+    task = dag.tasks[0]
+    if pool:
+        _maybe_display_run_warning(task)
+        # Use dummy run script for pool.
+        task.run = serve_constants.POOL_DUMMY_RUN_COMMAND
+
+    assert task.service is not None
+    if not pool and task.service.tls_credential is not None:
+        logger.warning('Updating TLS keyfile and certfile is not supported. '
+                       'Any updates to the keyfile and certfile will not take '
+                       'effect. To update TLS keyfile and certfile, please '
+                       'tear down the service and spin up a new one.')
+
+    prompt = None
+    if (service_record['status'] == serve_state.ServiceStatus.CONTROLLER_FAILED
+       ):
+        prompt = (f'{capnoun} {service_name!r} has a failed controller. '
+                  f'Please clean up the {noun} and try again.')
+    elif (service_record['status'] == serve_state.ServiceStatus.CONTROLLER_INIT
+         ):
+        prompt = (f'{capnoun} {service_name!r} is still initializing '
+                  'its controller. Please try again later.')
+    if prompt is not None:
+        with ux_utils.print_exception_no_traceback():
+            raise RuntimeError(prompt)
+
+    if not pool:
+        original_lb_policy = service_record['load_balancing_policy']
+        assert task.service is not None, 'Service section not found.'
+        if original_lb_policy != task.service.load_balancing_policy:
+            logger.warning(
+                f'{colorama.Fore.YELLOW}Current load balancing policy '
+                f'{original_lb_policy!r} is different from the new policy '
+                f'{task.service.load_balancing_policy!r}. Updating the load '
+                'balancing policy is not supported yet and it will be ignored. '
+                'The service will continue to use the current load balancing '
+                f'policy.{colorama.Style.RESET_ALL}')
+
+    with rich_utils.safe_status(
+            ux_utils.spinner_message(f'Initializing {noun}')):
+        controller_utils.maybe_translate_local_file_mounts_and_sync_up(
+            task, task_type='serve')
+
+    use_legacy = not handle.is_grpc_enabled_with_flag
+
+    if not use_legacy:
+        try:
+            current_version = serve_rpc_utils.RpcRunner.add_version(
+                handle, service_name)
+        except exceptions.SkyletMethodNotImplementedError:
+            use_legacy = True
+
+    if use_legacy:
+        code = serve_utils.ServeCodeGen.add_version(service_name)
+        returncode, version_string_payload, stderr = backend.run_on_head(
+            handle,
+            code,
+            require_outputs=True,
+            stream_logs=False,
+            separate_stderr=True)
+        try:
+            subprocess_utils.handle_returncode(returncode,
+                                               code,
+                                               'Failed to add version',
+                                               stderr,
+                                               stream_logs=True)
+        except exceptions.CommandError as e:
+            raise RuntimeError(e.error_msg) from e
+
+        version_string = serve_utils.load_version_string(version_string_payload)
+        try:
+            current_version = int(version_string)
+        except ValueError as e:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(f'Failed to parse version: {version_string}; '
+                                 f'Returncode: {returncode}') from e
+
+    with tempfile.NamedTemporaryFile(
+            prefix=f'{service_name}-v{current_version}',
+            mode='w') as service_file:
+        task_config = task.to_yaml_config()
+        yaml_utils.dump_yaml(service_file.name, task_config)
+        remote_task_yaml_path = serve_utils.generate_task_yaml_file_name(
+            service_name, current_version, expand_user=False)
+
+        with sky_logging.silent():
+            backend.sync_file_mounts(handle,
+                                     {remote_task_yaml_path: service_file.name},
+                                     storage_mounts=None)
+
+        use_legacy = not handle.is_grpc_enabled_with_flag
+
+        if not use_legacy:
+            try:
+                serve_rpc_utils.RpcRunner.update_service(
+                    handle, service_name, current_version, mode, pool)
+            except exceptions.SkyletMethodNotImplementedError:
+                use_legacy = True
+
+        if use_legacy:
+            code = serve_utils.ServeCodeGen.update_service(service_name,
+                                                           current_version,
+                                                           mode=mode.value,
+                                                           pool=pool)
+            returncode, _, stderr = backend.run_on_head(handle,
+                                                        code,
+                                                        require_outputs=True,
+                                                        stream_logs=False,
+                                                        separate_stderr=True)
+            try:
+                subprocess_utils.handle_returncode(returncode,
+                                                   code,
+                                                   f'Failed to update {noun}s',
+                                                   stderr,
+                                                   stream_logs=True)
+            except exceptions.CommandError as e:
+                raise RuntimeError(e.error_msg) from e
+
+    cmd = 'sky jobs pool status' if pool else 'sky serve status'
+    logger.info(
+        f'{colorama.Fore.GREEN}{capnoun} {service_name!r} update scheduled.'
+        f'{colorama.Style.RESET_ALL}\n'
+        f'Please use {ux_utils.BOLD}{cmd} {service_name} '
+        f'{ux_utils.RESET_BOLD}to check the latest status.')
+
+    if pool:
+        logs_cmd = f'`sky jobs pool logs {service_name} <worker_id>`'
+        unit_noun = 'Workers'
+
+    else:
+        logs_cmd = f'`sky serve logs {service_name} <replica_id>`'
+        unit_noun = 'Replicas'
+    logger.info(
+        ux_utils.finishing_message(
+            f'Successfully updated {noun} {service_name!r} '
+            f'to version {current_version}.',
+            follow_up_message=
+            f'\n{unit_noun} are updating, use {ux_utils.BOLD}{logs_cmd}'
+            f'{ux_utils.RESET_BOLD} to check their status.'))
+
+
+def apply(
+    task: 'task_lib.Task',
+    workers: Optional[int],
+    service_name: str,
+    mode: serve_utils.UpdateMode = serve_utils.DEFAULT_UPDATE_MODE,
+    pool: bool = False,
+) -> None:
+    """Applies the config to the service or pool."""
+    with filelock.FileLock(serve_utils.get_service_filelock_path(service_name)):
+        try:
+            controller_type = controller_utils.get_controller_for_pool(pool)
+            handle = backend_utils.is_controller_accessible(
+                controller=controller_type, stopped_message='')
+            backend = backend_utils.get_backend_from_handle(handle)
+            assert isinstance(backend, backends.CloudVmRayBackend)
+            service_record = _get_service_record(service_name, pool, handle,
+                                                 backend)
+            if service_record is not None:
+                return update(task, service_name, mode, pool, workers)
+        except exceptions.ClusterNotUpError:
+            pass
+        up(task, service_name, pool)
+
+
+def down(
+    service_names: Optional[Union[str, List[str]]] = None,
+    all: bool = False,  # pylint: disable=redefined-builtin
+    purge: bool = False,
+    pool: bool = False,
+) -> None:
+    """Tears down a service or pool."""
+    noun = 'pool' if pool else 'service'
+    if service_names is None:
+        service_names = []
+    if isinstance(service_names, str):
+        service_names = [service_names]
+    controller_type = controller_utils.get_controller_for_pool(pool)
+    handle = backend_utils.is_controller_accessible(
+        controller=controller_type,
+        stopped_message=f'All {noun}s should have terminated.')
+
+    service_names_str = ','.join(service_names)
+    if sum([bool(service_names), all]) != 1:
+        argument_str = (f'{noun}_names={service_names_str}'
+                        if service_names else '')
+        argument_str += ' all' if all else ''
+        raise ValueError(f'Can only specify one of {noun}_names or all. '
+                         f'Provided {argument_str!r}.')
+
+    service_names = None if all else service_names
+
+    try:
+        assert isinstance(handle, backends.CloudVmRayResourceHandle)
+        use_legacy = not handle.is_grpc_enabled_with_flag
+
+        if not use_legacy:
+            try:
+                stdout = serve_rpc_utils.RpcRunner.terminate_services(
+                    handle, service_names, purge, pool)
+            except exceptions.SkyletMethodNotImplementedError:
+                use_legacy = True
+
+        if use_legacy:
+            backend = backend_utils.get_backend_from_handle(handle)
+            assert isinstance(backend, backends.CloudVmRayBackend)
+            code = serve_utils.ServeCodeGen.terminate_services(
+                service_names, purge, pool)
+
+            returncode, stdout, _ = backend.run_on_head(handle,
+                                                        code,
+                                                        require_outputs=True,
+                                                        stream_logs=False)
+
+            subprocess_utils.handle_returncode(returncode, code,
+                                               f'Failed to terminate {noun}',
+                                               stdout)
+    except exceptions.FetchClusterInfoError as e:
+        raise RuntimeError(
+            'Failed to fetch controller IP. Please refresh controller status '
+            f'by `sky status -r {controller_type.value.cluster_name}` and try '
+            'again.') from e
+    except exceptions.CommandError as e:
+        raise RuntimeError(e.error_msg) from e
+    except grpc.RpcError as e:
+        raise RuntimeError(f'{e.details()} ({e.code()})') from e
+    except grpc.FutureTimeoutError as e:
+        raise RuntimeError('gRPC timed out') from e
+
+    logger.info(stdout)
+
+
+def status(
+    service_names: Optional[Union[str, List[str]]] = None,
+    pool: bool = False,
+) -> List[Dict[str, Any]]:
+    """Gets statuses of services or pools."""
+    noun = 'pool' if pool else 'service'
+    if service_names is not None:
+        if isinstance(service_names, str):
+            service_names = [service_names]
+
+    try:
+        backend_utils.check_network_connection()
+    except exceptions.NetworkError as e:
+        with ux_utils.print_exception_no_traceback():
+            raise RuntimeError(f'Failed to refresh {noun}s status '
+                               'due to network error.') from e
+
+    controller_type = controller_utils.get_controller_for_pool(pool)
+    handle = backend_utils.is_controller_accessible(
+        controller=controller_type,
+        stopped_message=controller_type.value.default_hint_if_non_existent.
+        replace('service', noun))
+
+    assert isinstance(handle, backends.CloudVmRayResourceHandle)
+    use_legacy = not handle.is_grpc_enabled_with_flag
+
+    if not use_legacy:
+        try:
+            service_records = serve_rpc_utils.RpcRunner.get_service_status(
+                handle, service_names, pool)
+        except exceptions.SkyletMethodNotImplementedError:
+            use_legacy = True
+
+    if use_legacy:
+        backend = backend_utils.get_backend_from_handle(handle)
+        assert isinstance(backend, backends.CloudVmRayBackend)
+
+        code = serve_utils.ServeCodeGen.get_service_status(service_names,
+                                                           pool=pool)
+        returncode, serve_status_payload, stderr = backend.run_on_head(
+            handle,
+            code,
+            require_outputs=True,
+            stream_logs=False,
+            separate_stderr=True)
+
+        try:
+            subprocess_utils.handle_returncode(returncode,
+                                               code,
+                                               f'Failed to fetch {noun}s',
+                                               stderr,
+                                               stream_logs=True)
+        except exceptions.CommandError as e:
+            raise RuntimeError(e.error_msg) from e
+
+        service_records = serve_utils.load_service_status(serve_status_payload)
+
+    # Get the endpoint for each service
+    for service_record in service_records:
+        service_record['endpoint'] = None
+        # Pool doesn't have an endpoint.
+        if pool:
+            continue
+        if service_record['load_balancer_port'] is not None:
+            try:
+                lb_port = service_record['load_balancer_port']
+                if not serve_utils.is_consolidation_mode(pool):
+                    endpoint = backend_utils.get_endpoints(
+                        cluster=common.SKY_SERVE_CONTROLLER_NAME,
+                        port=lb_port).get(lb_port, None)
+                else:
+                    endpoint = f'localhost:{lb_port}'
+            except exceptions.ClusterNotUpError:
+                pass
+            else:
+                protocol = ('https'
+                            if service_record['tls_encrypted'] else 'http')
+                if endpoint is not None:
+                    endpoint = endpoint.replace('https://',
+                                                '').replace('http://', '')
+                service_record['endpoint'] = f'{protocol}://{endpoint}'
+
+    return service_records
+
+
+ServiceComponentOrStr = Union[str, serve_utils.ServiceComponent]
+
+
+def tail_logs(
+    service_name: str,
+    *,
+    target: ServiceComponentOrStr,
+    replica_id: Optional[int] = None,
+    follow: bool = True,
+    tail: Optional[int] = None,
+    pool: bool = False,
+) -> None:
+    """Tail logs of a service or pool."""
+    if isinstance(target, str):
+        target = serve_utils.ServiceComponent(target)
+
+    if pool and target == serve_utils.ServiceComponent.LOAD_BALANCER:
+        raise ValueError(f'Target {target} is not supported for pool.')
+
+    if target == serve_utils.ServiceComponent.REPLICA:
+        if replica_id is None:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(
+                    '`replica_id` must be specified when using target=REPLICA.')
+    else:
+        if replica_id is not None:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError('`replica_id` must be None when using '
+                                 'target=CONTROLLER/LOAD_BALANCER.')
+
+    controller_type = controller_utils.get_controller_for_pool(pool)
+    handle = backend_utils.is_controller_accessible(
+        controller=controller_type,
+        stopped_message=controller_type.value.default_hint_if_non_existent)
+
+    backend = backend_utils.get_backend_from_handle(handle)
+    assert isinstance(backend, backends.CloudVmRayBackend), backend
+
+    if target != serve_utils.ServiceComponent.REPLICA:
+        code = serve_utils.ServeCodeGen.stream_serve_process_logs(
+            service_name,
+            stream_controller=(
+                target == serve_utils.ServiceComponent.CONTROLLER),
+            follow=follow,
+            tail=tail,
+            pool=pool)
+    else:
+        assert replica_id is not None, service_name
+        code = serve_utils.ServeCodeGen.stream_replica_logs(service_name,
+                                                            replica_id,
+                                                            follow,
+                                                            tail=tail,
+                                                            pool=pool)
+
+    # With the stdin=subprocess.DEVNULL, the ctrl-c will not directly
+    # kill the process, so we need to handle it manually here.
+    if threading.current_thread() is threading.main_thread():
+        signal.signal(signal.SIGINT, backend_utils.interrupt_handler)
+        signal.signal(signal.SIGTSTP, backend_utils.stop_handler)
+
+    # Refer to the notes in
+    # sky/backends/cloud_vm_ray_backend.py::CloudVmRayBackend::tail_logs.
+    backend.run_on_head(handle,
+                        code,
+                        stream_logs=True,
+                        process_stream=False,
+                        ssh_mode=command_runner.SshMode.INTERACTIVE)
+
+
+def _get_all_replica_targets(
+        service_name: str, backend: backends.CloudVmRayBackend,
+        handle: backends.CloudVmRayResourceHandle,
+        pool: bool) -> Set[serve_utils.ServiceComponentTarget]:
+    """Helper function to get targets for all live replicas."""
+    assert isinstance(handle, backends.CloudVmRayResourceHandle)
+    use_legacy = not handle.is_grpc_enabled_with_flag
+
+    if not use_legacy:
+        try:
+            service_records = serve_rpc_utils.RpcRunner.get_service_status(
+                handle, [service_name], pool)
+        except exceptions.SkyletMethodNotImplementedError:
+            use_legacy = True
+
+    if use_legacy:
+        code = serve_utils.ServeCodeGen.get_service_status([service_name],
+                                                           pool=pool)
+        returncode, serve_status_payload, stderr = backend.run_on_head(
+            handle,
+            code,
+            require_outputs=True,
+            stream_logs=False,
+            separate_stderr=True)
+
+        try:
+            subprocess_utils.handle_returncode(returncode,
+                                               code,
+                                               'Failed to fetch services',
+                                               stderr,
+                                               stream_logs=True)
+        except exceptions.CommandError as e:
+            raise RuntimeError(e.error_msg) from e
+
+        service_records = serve_utils.load_service_status(serve_status_payload)
+
+    if not service_records:
+        raise ValueError(f'Service {service_name!r} not found.')
+    assert len(service_records) == 1
+    service_record = service_records[0]
+
+    return {
+        serve_utils.ServiceComponentTarget(serve_utils.ServiceComponent.REPLICA,
+                                           replica_info['replica_id'])
+        for replica_info in service_record['replica_info']
+    }
+
+
+def sync_down_logs(
+    service_name: str,
+    *,
+    local_dir: str,
+    targets: Union[ServiceComponentOrStr, List[ServiceComponentOrStr],
+                   None] = None,
+    replica_ids: Optional[List[int]] = None,
+    tail: Optional[int] = None,
+    pool: bool = False,
+) -> str:
+    """Sync down logs of a service or pool."""
+    noun = 'pool' if pool else 'service'
+    repnoun = 'worker' if pool else 'replica'
+    caprepnoun = repnoun.capitalize()
+
+    # Step 0) get the controller handle
+    with rich_utils.safe_status(
+            ux_utils.spinner_message(f'Checking {noun} status...')):
+        controller_type = controller_utils.get_controller_for_pool(pool)
+        handle = backend_utils.is_controller_accessible(
+            controller=controller_type,
+            stopped_message=controller_type.value.default_hint_if_non_existent)
+        backend: backends.CloudVmRayBackend = (
+            backend_utils.get_backend_from_handle(handle))
+
+    requested_components: Set[serve_utils.ServiceComponent] = set()
+    if not targets:
+        # No targets specified -> request all components
+        requested_components = {
+            serve_utils.ServiceComponent.CONTROLLER,
+            serve_utils.ServiceComponent.LOAD_BALANCER,
+            serve_utils.ServiceComponent.REPLICA
+        }
+    else:
+        # Parse provided targets
+        if isinstance(targets, (str, serve_utils.ServiceComponent)):
+            requested_components = {serve_utils.ServiceComponent(targets)}
+        else:  # list
+            requested_components = {
+                serve_utils.ServiceComponent(t) for t in targets
+            }
+
+    normalized_targets: Set[serve_utils.ServiceComponentTarget] = set()
+    if serve_utils.ServiceComponent.CONTROLLER in requested_components:
+        normalized_targets.add(
+            serve_utils.ServiceComponentTarget(
+                serve_utils.ServiceComponent.CONTROLLER))
+    if serve_utils.ServiceComponent.LOAD_BALANCER in requested_components:
+        normalized_targets.add(
+            serve_utils.ServiceComponentTarget(
+                serve_utils.ServiceComponent.LOAD_BALANCER))
+    if serve_utils.ServiceComponent.REPLICA in requested_components:
+        with rich_utils.safe_status(
+                ux_utils.spinner_message(f'Getting live {repnoun} infos...')):
+            replica_targets = _get_all_replica_targets(service_name, backend,
+                                                       handle, pool)
+        if not replica_ids:
+            # Replica target requested but no specific IDs
+            # -> Get all replica logs
+            normalized_targets.update(replica_targets)
+        else:
+            # Replica target requested with specific IDs
+            requested_replica_targets = [
+                serve_utils.ServiceComponentTarget(
+                    serve_utils.ServiceComponent.REPLICA, rid)
+                for rid in replica_ids
+            ]
+            for target in requested_replica_targets:
+                if target not in replica_targets:
+                    logger.warning(f'{caprepnoun} ID {target.replica_id} not '
+                                   f'found for {service_name}. Skipping...')
+                else:
+                    normalized_targets.add(target)
+
+    def sync_down_logs_by_target(target: serve_utils.ServiceComponentTarget):
+        component = target.component
+        # We need to set one side of the pipe to a logs stream, and the other
+        # side to a file.
+        log_path = str(pathlib.Path(local_dir) / f'{target}.log')
+        stream_logs_code: str
+
+        if component == serve_utils.ServiceComponent.CONTROLLER:
+            stream_logs_code = (
+                serve_utils.ServeCodeGen.stream_serve_process_logs(
+                    service_name,
+                    stream_controller=True,
+                    follow=False,
+                    tail=tail,
+                    pool=pool))
+        elif component == serve_utils.ServiceComponent.LOAD_BALANCER:
+            stream_logs_code = (
+                serve_utils.ServeCodeGen.stream_serve_process_logs(
+                    service_name,
+                    stream_controller=False,
+                    follow=False,
+                    tail=tail,
+                    pool=pool))
+        elif component == serve_utils.ServiceComponent.REPLICA:
+            replica_id = target.replica_id
+            assert replica_id is not None, service_name
+            stream_logs_code = serve_utils.ServeCodeGen.stream_replica_logs(
+                service_name, replica_id, follow=False, tail=tail, pool=pool)
+        else:
+            assert False, component
+
+        # Refer to the notes in
+        # sky/backends/cloud_vm_ray_backend.py::CloudVmRayBackend::tail_logs.
+        backend.run_on_head(handle,
+                            stream_logs_code,
+                            stream_logs=False,
+                            process_stream=False,
+                            ssh_mode=command_runner.SshMode.INTERACTIVE,
+                            log_path=log_path)
+
+    subprocess_utils.run_in_parallel(sync_down_logs_by_target,
+                                     list(normalized_targets))
+
+    return local_dir