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

sky/serve/serve_utils.py

@@ -1,6 +1,7 @@
 """User interface with the SkyServe."""
 import base64
 import collections
+import dataclasses
 import enum
 import os
 import pathlib
@@ -23,15 +24,15 @@ import requests
 from sky import backends
 from sky import exceptions
 from sky import global_user_state
-from sky import status_lib
-from sky.backends import backend_utils
 from sky.serve import constants
 from sky.serve import serve_state
 from sky.skylet import constants as skylet_constants
 from sky.skylet import job_lib
 from sky.utils import common_utils
 from sky.utils import log_utils
+from sky.utils import message_utils
 from sky.utils import resources_utils
+from sky.utils import status_lib
 from sky.utils import ux_utils
 
 if typing.TYPE_CHECKING:
@@ -39,15 +40,19 @@ if typing.TYPE_CHECKING:
 
     from sky.serve import replica_managers
 
-SKY_SERVE_CONTROLLER_NAME: str = (
-    f'sky-serve-controller-{common_utils.get_user_hash()}')
 _SYSTEM_MEMORY_GB = psutil.virtual_memory().total // (1024**3)
 NUM_SERVICE_THRESHOLD = (_SYSTEM_MEMORY_GB //
                          constants.CONTROLLER_MEMORY_USAGE_GB)
 _CONTROLLER_URL = 'http://localhost:{CONTROLLER_PORT}'
 
-_SKYPILOT_PROVISION_LOG_PATTERN = r'.*tail -n100 -f (.*provision\.log).*'
-_SKYPILOT_LOG_PATTERN = r'.*tail -n100 -f (.*\.log).*'
+# NOTE(dev): We assume log paths are either in ~/sky_logs/... or ~/.sky/...
+# and always appear after a space. Be careful when changing UX as this
+# assumption is used to expand some log files while ignoring others.
+_SKYPILOT_LOG_DIRS = r'~/(sky_logs|\.sky)'
+_SKYPILOT_PROVISION_LOG_PATTERN = (
+    fr'.* ({_SKYPILOT_LOG_DIRS}/.*provision\.log)')
+_SKYPILOT_LOG_PATTERN = fr'.* ({_SKYPILOT_LOG_DIRS}/.*\.log)'
+
 # TODO(tian): Find all existing replica id and print here.
 _FAILED_TO_FIND_REPLICA_MSG = (
     f'{colorama.Fore.RED}Failed to find replica '
@@ -86,6 +91,19 @@ class UpdateMode(enum.Enum):
     BLUE_GREEN = 'blue_green'
 
 
+@dataclasses.dataclass
+class TLSCredential:
+    """TLS credential for the service."""
+    keyfile: str
+    certfile: str
+
+    def dump_uvicorn_kwargs(self) -> Dict[str, str]:
+        return {
+            'ssl_keyfile': os.path.expanduser(self.keyfile),
+            'ssl_certfile': os.path.expanduser(self.certfile),
+        }
+
+
 DEFAULT_UPDATE_MODE = UpdateMode.ROLLING
 
 _SIGNAL_TO_ERROR = {
@@ -104,7 +122,7 @@ ValueType = TypeVar('ValueType')
 class ThreadSafeDict(Generic[KeyType, ValueType]):
     """A thread-safe dict."""
 
-    def __init__(self, *args, **kwargs) -> None:
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         self._dict: Dict[KeyType, ValueType] = dict(*args, **kwargs)
         self._lock = threading.Lock()
 
@@ -237,6 +255,18 @@ def generate_replica_log_file_name(service_name: str, replica_id: int) -> str:
     return os.path.join(dir_name, f'replica_{replica_id}.log')
 
 
+def generate_remote_tls_keyfile_name(service_name: str) -> str:
+    dir_name = generate_remote_service_dir_name(service_name)
+    # Don't expand here since it is used for remote machine.
+    return os.path.join(dir_name, 'tls_keyfile')
+
+
+def generate_remote_tls_certfile_name(service_name: str) -> str:
+    dir_name = generate_remote_service_dir_name(service_name)
+    # Don't expand here since it is used for remote machine.
+    return os.path.join(dir_name, 'tls_certfile')
+
+
 def generate_replica_cluster_name(service_name: str, replica_id: int) -> str:
     return f'{service_name}-{replica_id}'
 
@@ -246,9 +276,11 @@ def set_service_status_and_active_versions_from_replica(
         update_mode: UpdateMode) -> None:
     record = serve_state.get_service_from_name(service_name)
     if record is None:
-        raise ValueError('The service is up-ed in an old version and does not '
-                         'support update. Please `sky serve down` '
-                         'it first and relaunch the service.')
+        with ux_utils.print_exception_no_traceback():
+            raise ValueError(
+                'The service is up-ed in an old version and does not '
+                'support update. Please `sky serve down` '
+                'it first and relaunch the service.')
     if record['status'] == serve_state.ServiceStatus.SHUTTING_DOWN:
         # When the service is shutting down, there is a period of time which the
         # controller still responds to the request, and the replica is not
@@ -289,7 +321,8 @@ def update_service_status() -> None:
 def update_service_encoded(service_name: str, version: int, mode: str) -> str:
     service_status = _get_service_status(service_name)
     if service_status is None:
-        raise ValueError(f'Service {service_name!r} does not exist.')
+        with ux_utils.print_exception_no_traceback():
+            raise ValueError(f'Service {service_name!r} does not exist.')
     controller_port = service_status['controller_port']
     resp = requests.post(
         _CONTROLLER_URL.format(CONTROLLER_PORT=controller_port) +
@@ -299,14 +332,54 @@ def update_service_encoded(service_name: str, version: int, mode: str) -> str:
             'mode': mode,
         })
     if resp.status_code == 404:
-        raise ValueError('The service is up-ed in an old version and does not '
-                         'support update. Please `sky serve down` '
-                         'it first and relaunch the service. ')
+        with ux_utils.print_exception_no_traceback():
+            raise ValueError(
+                'The service is up-ed in an old version and does not '
+                'support update. Please `sky serve down` '
+                'it first and relaunch the service. ')
+    elif resp.status_code == 400:
+        with ux_utils.print_exception_no_traceback():
+            raise ValueError(f'Client error during service update: {resp.text}')
+    elif resp.status_code == 500:
+        with ux_utils.print_exception_no_traceback():
+            raise RuntimeError(
+                f'Server error during service update: {resp.text}')
     elif resp.status_code != 200:
-        raise ValueError(f'Failed to update service: {resp.text}')
+        with ux_utils.print_exception_no_traceback():
+            raise ValueError(f'Failed to update service: {resp.text}')
 
     service_msg = resp.json()['message']
-    return common_utils.encode_payload(service_msg)
+    return message_utils.encode_payload(service_msg)
+
+
+def terminate_replica(service_name: str, replica_id: int, purge: bool) -> str:
+    service_status = _get_service_status(service_name)
+    if service_status is None:
+        with ux_utils.print_exception_no_traceback():
+            raise ValueError(f'Service {service_name!r} does not exist.')
+    replica_info = serve_state.get_replica_info_from_id(service_name,
+                                                        replica_id)
+    if replica_info is None:
+        with ux_utils.print_exception_no_traceback():
+            raise ValueError(
+                f'Replica {replica_id} for service {service_name} does not '
+                'exist.')
+
+    controller_port = service_status['controller_port']
+    resp = requests.post(
+        _CONTROLLER_URL.format(CONTROLLER_PORT=controller_port) +
+        '/controller/terminate_replica',
+        json={
+            'replica_id': replica_id,
+            'purge': purge,
+        })
+
+    message: str = resp.json()['message']
+    if resp.status_code != 200:
+        with ux_utils.print_exception_no_traceback():
+            raise ValueError(f'Failed to terminate replica {replica_id} '
+                             f'in {service_name}. Reason:\n{message}')
+    return message
 
 
 def _get_service_status(
@@ -334,7 +407,7 @@ def _get_service_status(
 
 
 def get_service_status_encoded(service_names: Optional[List[str]]) -> str:
-    service_statuses = []
+    service_statuses: List[Dict[str, str]] = []
     if service_names is None:
         # Get all service names
         service_names = serve_state.get_glob_service_names(None)
@@ -346,13 +419,28 @@ def get_service_status_encoded(service_names: Optional[List[str]]) -> str:
             k: base64.b64encode(pickle.dumps(v)).decode('utf-8')
             for k, v in service_status.items()
         })
-    return common_utils.encode_payload(service_statuses)
+    # We have to use payload_type here to avoid the issue of
+    # message_utils.decode_payload() not being able to correctly decode the
+    # message with <sky-payload> tags.
+    return message_utils.encode_payload(service_statuses,
+                                        payload_type='service_status')
 
 
 def load_service_status(payload: str) -> List[Dict[str, Any]]:
-    service_statuses_encoded = common_utils.decode_payload(payload)
-    service_statuses = []
+    try:
+        service_statuses_encoded = message_utils.decode_payload(
+            payload, payload_type='service_status')
+    except ValueError as e:
+        if 'Invalid payload string' in str(e):
+            # Backward compatibility for serve controller started before #4660
+            # where the payload type is not added.
+            service_statuses_encoded = message_utils.decode_payload(payload)
+        else:
+            raise
+    service_statuses: List[Dict[str, Any]] = []
     for service_status in service_statuses_encoded:
+        if not isinstance(service_status, dict):
+            raise ValueError(f'Invalid service status: {service_status}')
         service_statuses.append({
             k: pickle.loads(base64.b64decode(v))
             for k, v in service_status.items()
@@ -362,16 +450,16 @@ def load_service_status(payload: str) -> List[Dict[str, Any]]:
 
 def add_version_encoded(service_name: str) -> str:
     new_version = serve_state.add_version(service_name)
-    return common_utils.encode_payload(new_version)
+    return message_utils.encode_payload(new_version)
 
 
 def load_version_string(payload: str) -> str:
-    return common_utils.decode_payload(payload)
+    return message_utils.decode_payload(payload)
 
 
 def _terminate_failed_services(
         service_name: str,
-        service_status: serve_state.ServiceStatus) -> Optional[str]:
+        service_status: Optional[serve_state.ServiceStatus]) -> Optional[str]:
     """Terminate service in failed status.
 
     Services included in ServiceStatus.failed_statuses() do not have an
@@ -383,7 +471,7 @@ def _terminate_failed_services(
         A message indicating potential resource leak (if any). If no
         resource leak is detected, return None.
     """
-    remaining_replica_clusters = []
+    remaining_replica_clusters: List[str] = []
     # The controller should have already attempted to terminate those
     # replicas, so we don't need to try again here.
     for replica_info in serve_state.get_replica_infos(service_name):
@@ -397,6 +485,7 @@ def _terminate_failed_services(
         generate_remote_service_dir_name(service_name))
     shutil.rmtree(service_dir)
     serve_state.remove_service(service_name)
+    serve_state.delete_all_versions(service_name)
 
     if not remaining_replica_clusters:
         return None
@@ -409,26 +498,35 @@ def _terminate_failed_services(
 
 def terminate_services(service_names: Optional[List[str]], purge: bool) -> str:
     service_names = serve_state.get_glob_service_names(service_names)
-    terminated_service_names = []
-    messages = []
+    terminated_service_names: List[str] = []
+    messages: List[str] = []
     for service_name in service_names:
         service_status = _get_service_status(service_name,
                                              with_replica_info=False)
-        assert service_status is not None
-        if service_status['status'] == serve_state.ServiceStatus.SHUTTING_DOWN:
+        if (service_status is not None and service_status['status']
+                == serve_state.ServiceStatus.SHUTTING_DOWN):
             # Already scheduled to be terminated.
             continue
-        if (service_status['status']
+        # If the `services` and `version_specs` table are not aligned, it might
+        # result in a None service status. In this case, the controller process
+        # is not functioning as well and we should also use the
+        # `_terminate_failed_services` function to clean up the service.
+        # This is a safeguard for a rare case, that is accidentally abort
+        # between `serve_state.add_service` and
+        # `serve_state.add_or_update_version` in service.py.
+        if (service_status is None or service_status['status']
                 in serve_state.ServiceStatus.failed_statuses()):
+            failed_status = (service_status['status']
+                             if service_status is not None else None)
             if purge:
                 message = _terminate_failed_services(service_name,
-                                                     service_status['status'])
+                                                     failed_status)
                 if message is not None:
                     messages.append(message)
             else:
                 messages.append(
                     f'{colorama.Fore.YELLOW}Service {service_name!r} is in '
-                    f'failed status ({service_status["status"]}). Skipping '
+                    f'failed status ({failed_status}). Skipping '
                     'its termination as it could lead to a resource leak. '
                     f'(Use `sky serve down {service_name} --purge` to '
                     'forcefully terminate the service.)'
@@ -447,7 +545,7 @@ def terminate_services(service_names: Optional[List[str]], purge: bool) -> str:
                     f.write(UserSignal.TERMINATE.value)
                     f.flush()
         terminated_service_names.append(f'{service_name!r}')
-    if len(terminated_service_names) == 0:
+    if not terminated_service_names:
         messages.append('No service to terminate.')
     else:
         identity_str = f'Service {terminated_service_names[0]} is'
@@ -472,7 +570,31 @@ def wait_service_registration(service_name: str, job_id: int) -> str:
         Encoded load balancer port assigned to the service.
     """
     start_time = time.time()
+    setup_completed = False
     while True:
+        job_status = job_lib.get_status(job_id)
+        if job_status is None or job_status < job_lib.JobStatus.RUNNING:
+            # Wait for the controller process to finish setting up. It can be
+            # slow if a lot cloud dependencies are being installed.
+            if (time.time() - start_time >
+                    constants.CONTROLLER_SETUP_TIMEOUT_SECONDS):
+                with ux_utils.print_exception_no_traceback():
+                    raise RuntimeError(
+                        f'Failed to start the controller '
+                        f'process for the service {service_name!r} '
+                        f'within '
+                        f'{constants.CONTROLLER_SETUP_TIMEOUT_SECONDS} seconds.'
+                    )
+            # No need to check the service status as the controller process
+            # is still setting up.
+            time.sleep(1)
+            continue
+
+        if not setup_completed:
+            setup_completed = True
+            # Reset the start time to wait for the service to be registered.
+            start_time = time.time()
+
         record = serve_state.get_service_from_name(service_name)
         if record is not None:
             if job_id != record['controller_job_id']:
@@ -480,12 +602,11 @@ def wait_service_registration(service_name: str, job_id: int) -> str:
                     raise ValueError(
                         f'The service {service_name!r} is already running. '
                         'Please specify a different name for your service. '
-                        'To update an existing service, run: `sky serve down` '
-                        'and then `sky serve up` again (in-place update will '
-                        'be supported in the future).')
+                        'To update an existing service, run: sky serve update '
+                        f'{service_name} <new-service-yaml>')
             lb_port = record['load_balancer_port']
             if lb_port is not None:
-                return common_utils.encode_payload(lb_port)
+                return message_utils.encode_payload(lb_port)
         elif len(serve_state.get_services()) >= NUM_SERVICE_THRESHOLD:
             with ux_utils.print_exception_no_traceback():
                 raise RuntimeError('Max number of services reached. '
@@ -508,7 +629,7 @@ def wait_service_registration(service_name: str, job_id: int) -> str:
 
 
 def load_service_initialization_result(payload: str) -> int:
-    return common_utils.decode_payload(payload)
+    return message_utils.decode_payload(payload)
 
 
 def check_service_status_healthy(service_name: str) -> Optional[str]:
@@ -539,16 +660,27 @@ def get_latest_version_with_min_replicas(
     return active_versions[-1] if active_versions else None
 
 
-def _follow_replica_logs(
-        file: TextIO,
-        cluster_name: str,
-        *,
-        finish_stream: Callable[[], bool],
-        exit_if_stream_end: bool = False,
-        no_new_content_timeout: Optional[int] = None) -> Iterator[str]:
-    line = ''
-    log_file = None
-    no_new_content_cnt = 0
+def _follow_logs_with_provision_expanding(
+    file: TextIO,
+    cluster_name: str,
+    *,
+    should_stop: Callable[[], bool],
+    stop_on_eof: bool = False,
+    idle_timeout_seconds: Optional[int] = None,
+) -> Iterator[str]:
+    """Follows logs and expands any provision.log references found.
+
+    Args:
+        file: Log file to read from.
+        cluster_name: Name of the cluster being launched.
+        should_stop: Callback that returns True when streaming should stop.
+        stop_on_eof: If True, stop when reaching end of file.
+        idle_timeout_seconds: If set, stop after these many seconds without
+            new content.
+
+    Yields:
+        Log lines, including expanded content from referenced provision logs.
+    """
 
     def cluster_is_up() -> bool:
         cluster_record = global_user_state.get_cluster_from_name(cluster_name)
@@ -556,51 +688,51 @@ def _follow_replica_logs(
             return False
         return cluster_record['status'] == status_lib.ClusterStatus.UP
 
-    while True:
-        tmp = file.readline()
-        if tmp is not None and tmp != '':
-            no_new_content_cnt = 0
-            line += tmp
-            if '\n' in line or '\r' in line:
-                # Tailing detailed progress for user. All logs in skypilot is
-                # of format `To view detailed progress: tail -n100 -f *.log`.
-                x = re.match(_SKYPILOT_PROVISION_LOG_PATTERN, line)
-                if x is not None:
-                    log_file = os.path.expanduser(x.group(1))
-                elif re.match(_SKYPILOT_LOG_PATTERN, line) is None:
-                    # Not print other logs (file sync logs) since we lack
-                    # utility to determine when these log files are finished
-                    # writing.
-                    # TODO(tian): Not skip these logs since there are small
-                    # chance that error will happen in file sync. Need to find
-                    # a better way to do this.
-                    yield line
-                    # Output next line first since it indicates the process is
-                    # starting. For our launching logs, it's always:
-                    # Launching on <cloud> <region> (<zone>)
-                    if log_file is not None:
-                        with open(log_file, 'r', newline='',
-                                  encoding='utf-8') as f:
-                            # We still exit if more than 10 seconds without new
-                            # content to avoid any internal bug that causes
-                            # the launch failed and cluster status remains INIT.
-                            for l in _follow_replica_logs(
-                                    f,
-                                    cluster_name,
-                                    finish_stream=cluster_is_up,
-                                    exit_if_stream_end=exit_if_stream_end,
-                                    no_new_content_timeout=10):
-                                yield l
-                        log_file = None
-                line = ''
-        else:
-            if exit_if_stream_end or finish_stream():
-                break
-            if no_new_content_timeout is not None:
-                if no_new_content_cnt >= no_new_content_timeout:
-                    break
-                no_new_content_cnt += 1
-            time.sleep(1)
+    def process_line(line: str) -> Iterator[str]:
+        # The line might be directing users to view logs, like
+        # `✓ Cluster launched: new-http.  View logs at: *.log`
+        # We should tail the detailed logs for user.
+        provision_log_prompt = re.match(_SKYPILOT_PROVISION_LOG_PATTERN, line)
+        log_prompt = re.match(_SKYPILOT_LOG_PATTERN, line)
+
+        if provision_log_prompt is not None:
+            nested_log_path = os.path.expanduser(provision_log_prompt.group(1))
+
+            try:
+                with open(nested_log_path, 'r', newline='',
+                          encoding='utf-8') as f:
+                    # We still exit if more than 10 seconds without new content
+                    # to avoid any internal bug that causes the launch to fail
+                    # while cluster status remains INIT.
+                    yield from log_utils.follow_logs(f,
+                                                     should_stop=cluster_is_up,
+                                                     stop_on_eof=stop_on_eof,
+                                                     idle_timeout_seconds=10)
+            except FileNotFoundError:
+                yield line
+
+                yield (f'{colorama.Fore.YELLOW}{colorama.Style.BRIGHT}'
+                       f'Try to expand log file {nested_log_path} but not '
+                       f'found. Skipping...{colorama.Style.RESET_ALL}')
+                pass
+            return
+
+        if log_prompt is not None:
+            # Now we skip other logs (file sync logs) since we lack
+            # utility to determine when these log files are finished
+            # writing.
+            # TODO(tian): We should not skip these logs since there are
+            # small chance that error will happen in file sync. Need to
+            # find a better way to do this.
+            return
+
+        yield line
+
+    return log_utils.follow_logs(file,
+                                 should_stop=should_stop,
+                                 stop_on_eof=stop_on_eof,
+                                 process_line=process_line,
+                                 idle_timeout_seconds=idle_timeout_seconds)
 
 
 def stream_replica_logs(service_name: str, replica_id: int,
@@ -631,17 +763,21 @@ def stream_replica_logs(service_name: str, replica_id: int,
         for info in replica_info:
             if info.replica_id == replica_id:
                 return info.status
-        raise ValueError(
-            _FAILED_TO_FIND_REPLICA_MSG.format(replica_id=replica_id))
+        with ux_utils.print_exception_no_traceback():
+            raise ValueError(
+                _FAILED_TO_FIND_REPLICA_MSG.format(replica_id=replica_id))
 
-    finish_stream = (
+    replica_provisioned = (
         lambda: _get_replica_status() != serve_state.ReplicaStatus.PROVISIONING)
     with open(launch_log_file_name, 'r', newline='', encoding='utf-8') as f:
-        for line in _follow_replica_logs(f,
-                                         replica_cluster_name,
-                                         finish_stream=finish_stream,
-                                         exit_if_stream_end=not follow):
+        for line in _follow_logs_with_provision_expanding(
+                f,
+                replica_cluster_name,
+                should_stop=replica_provisioned,
+                stop_on_eof=not follow,
+        ):
             print(line, end='', flush=True)
+
     if (not follow and
             _get_replica_status() == serve_state.ReplicaStatus.PROVISIONING):
         # Early exit if not following the logs.
@@ -666,22 +802,6 @@ def stream_replica_logs(service_name: str, replica_id: int,
     return ''
 
 
-def _follow_logs(file: TextIO, *, finish_stream: Callable[[], bool],
-                 exit_if_stream_end: bool) -> Iterator[str]:
-    line = ''
-    while True:
-        tmp = file.readline()
-        if tmp is not None and tmp != '':
-            line += tmp
-            if '\n' in line or '\r' in line:
-                yield line
-                line = ''
-        else:
-            if exit_if_stream_end or finish_stream():
-                break
-            time.sleep(1)
-
-
 def stream_serve_process_logs(service_name: str, stream_controller: bool,
                               follow: bool) -> str:
     msg = check_service_status_healthy(service_name)
@@ -700,9 +820,11 @@ def stream_serve_process_logs(service_name: str, stream_controller: bool,
 
     with open(os.path.expanduser(log_file), 'r', newline='',
               encoding='utf-8') as f:
-        for line in _follow_logs(f,
-                                 finish_stream=_service_is_terminal,
-                                 exit_if_stream_end=not follow):
+        for line in log_utils.follow_logs(
+                f,
+                should_stop=_service_is_terminal,
+                stop_on_eof=not follow,
+        ):
             print(line, end='', flush=True)
     return ''
 
@@ -721,28 +843,6 @@ def _get_replicas(service_record: Dict[str, Any]) -> str:
     return f'{ready_replica_num}/{total_replica_num}'
 
 
-def get_endpoint(service_record: Dict[str, Any]) -> str:
-    # Don't use backend_utils.is_controller_up since it is too slow.
-    handle = global_user_state.get_handle_from_cluster_name(
-        SKY_SERVE_CONTROLLER_NAME)
-    assert isinstance(handle, backends.CloudVmRayResourceHandle)
-    if handle is None:
-        return '-'
-    load_balancer_port = service_record['load_balancer_port']
-    if load_balancer_port is None:
-        return '-'
-    try:
-        endpoint = backend_utils.get_endpoints(handle.cluster_name,
-                                               load_balancer_port).get(
-                                                   load_balancer_port, None)
-    except exceptions.ClusterNotUpError:
-        return '-'
-    if endpoint is None:
-        return '-'
-    assert isinstance(endpoint, str), endpoint
-    return endpoint
-
-
 def format_service_table(service_records: List[Dict[str, Any]],
                          show_all: bool) -> str:
     if not service_records:
@@ -752,10 +852,12 @@ def format_service_table(service_records: List[Dict[str, Any]],
         'NAME', 'VERSION', 'UPTIME', 'STATUS', 'REPLICAS', 'ENDPOINT'
     ]
     if show_all:
-        service_columns.extend(['POLICY', 'REQUESTED_RESOURCES'])
+        service_columns.extend([
+            'AUTOSCALING_POLICY', 'LOAD_BALANCING_POLICY', 'REQUESTED_RESOURCES'
+        ])
     service_table = log_utils.create_table(service_columns)
 
-    replica_infos = []
+    replica_infos: List[Dict[str, Any]] = []
     for record in service_records:
         for replica in record['replica_info']:
             replica['service_name'] = record['name']
@@ -770,14 +872,12 @@ def format_service_table(service_records: List[Dict[str, Any]],
         service_status = record['status']
         status_str = service_status.colored_str()
         replicas = _get_replicas(record)
-        endpoint = get_endpoint(record)
+        endpoint = record['endpoint']
+        if endpoint is None:
+            endpoint = '-'
         policy = record['policy']
-        # TODO(tian): Backward compatibility.
-        # Remove `requested_resources` field after 2 minor release, 0.6.0.
-        if record.get('requested_resources_str') is None:
-            requested_resources_str = str(record['requested_resources'])
-        else:
-            requested_resources_str = record['requested_resources_str']
+        requested_resources_str = record['requested_resources_str']
+        load_balancing_policy = record['load_balancing_policy']
 
         service_values = [
             service_name,
@@ -788,7 +888,8 @@ def format_service_table(service_records: List[Dict[str, Any]],
             endpoint,
         ]
         if show_all:
-            service_values.extend([policy, requested_resources_str])
+            service_values.extend(
+                [policy, load_balancing_policy, requested_resources_str])
         service_table.add_row(service_values)
 
     replica_table = _format_replica_table(replica_infos, show_all)
@@ -830,7 +931,8 @@ def _format_replica_table(replica_records: List[Dict[str, Any]],
         region = '-'
         zone = '-'
 
-        replica_handle: 'backends.CloudVmRayResourceHandle' = record['handle']
+        replica_handle: Optional['backends.CloudVmRayResourceHandle'] = record[
+            'handle']
         if replica_handle is not None:
             resources_str = resources_utils.get_readable_resources_repr(
                 replica_handle, simplify=not show_all)
@@ -902,6 +1004,18 @@ class ServeCodeGen:
         ]
         return cls._build(code)
 
+    @classmethod
+    def terminate_replica(cls, service_name: str, replica_id: int,
+                          purge: bool) -> str:
+        code = [
+            f'(lambda: print(serve_utils.terminate_replica({service_name!r}, '
+            f'{replica_id}, {purge}), end="", flush=True) '
+            'if getattr(constants, "SERVE_VERSION", 0) >= 2 else '
+            f'exec("raise RuntimeError('
+            f'{constants.TERMINATE_REPLICA_VERSION_MISMATCH_ERROR!r})"))()'
+        ]
+        return cls._build(code)
+
     @classmethod
     def wait_service_registration(cls, service_name: str, job_id: int) -> str:
         code = [
@@ -933,21 +1047,18 @@ class ServeCodeGen:
     def _build(cls, code: List[str]) -> str:
         code = cls._PREFIX + code
         generated_code = '; '.join(code)
-        return (f'{skylet_constants.SKY_PYTHON_CMD} '
+        # Use the local user id to make sure the operation goes to the correct
+        # user.
+        return (f'export {skylet_constants.USER_ID_ENV_VAR}='
+                f'"{common_utils.get_user_hash()}"; '
+                f'{skylet_constants.SKY_PYTHON_CMD} '
                 f'-u -c {shlex.quote(generated_code)}')
 
     @classmethod
     def update_service(cls, service_name: str, version: int, mode: str) -> str:
         code = [
-            # Backward compatibility for old serve version on the remote
-            # machine. The `mode` argument was added in #3249, and if the remote
-            # machine has an old SkyPilot version before that, we need to avoid
-            # passing the `mode` argument to the job_lib functions.
-            # TODO(zhwu): Remove this in 0.7.0 release.
-            f'mode_kwargs = {{"mode": {mode!r}}} '
-            'if getattr(constants, "SERVE_VERSION", 0) >= 1 else {}',
             f'msg = serve_utils.update_service_encoded({service_name!r}, '
-            f'{version}, **mode_kwargs)',
+            f'{version}, mode={mode!r})',
             'print(msg, end="", flush=True)',
         ]
         return cls._build(code)