skypilot-nightly 1.0.0.dev20250802__py3-none-any.whl → 1.0.0.dev20250806__py3-none-any.whl

sky/serve/client/sdk_async.py

@@ -0,0 +1,130 @@
+"""Async SDK for SkyServe."""
+import typing
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+from sky.client import sdk_async
+from sky.serve.client import sdk
+from sky.usage import usage_lib
+from sky.utils import context_utils
+
+if typing.TYPE_CHECKING:
+    import io
+
+    import sky
+    from sky.serve import serve_utils
+
+
+@usage_lib.entrypoint
+async def up(
+    task: Union['sky.Task', 'sky.Dag'],
+    service_name: str,
+    # Internal only:
+    # pylint: disable=invalid-name
+    _need_confirmation: bool = False,
+    stream_logs: Optional[
+        sdk_async.StreamConfig] = sdk_async.DEFAULT_STREAM_CONFIG
+) -> Tuple[str, str]:
+    """Async version of up() that spins up a service."""
+    request_id = await context_utils.to_thread(sdk.up, task, service_name,
+                                               _need_confirmation)
+    if stream_logs is not None:
+        return await sdk_async._stream_and_get(request_id, stream_logs)  # pylint: disable=protected-access
+    else:
+        return await sdk_async.get(request_id)
+
+
+@usage_lib.entrypoint
+async 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,
+    stream_logs: Optional[
+        sdk_async.StreamConfig] = sdk_async.DEFAULT_STREAM_CONFIG
+) -> None:
+    """Async version of update() that updates an existing service."""
+    request_id = await context_utils.to_thread(sdk.update, task, service_name,
+                                               mode, _need_confirmation)
+    if stream_logs is not None:
+        return await sdk_async._stream_and_get(request_id, stream_logs)  # pylint: disable=protected-access
+    else:
+        return await sdk_async.get(request_id)
+
+
+@usage_lib.entrypoint
+async def down(
+    service_names: Optional[Union[str, List[str]]],
+    all: bool = False,  # pylint: disable=redefined-builtin
+    purge: bool = False,
+    stream_logs: Optional[
+        sdk_async.StreamConfig] = sdk_async.DEFAULT_STREAM_CONFIG
+) -> None:
+    """Async version of down() that tears down a service."""
+    request_id = await context_utils.to_thread(sdk.down, service_names, all,
+                                               purge)
+    if stream_logs is not None:
+        return await sdk_async._stream_and_get(request_id, stream_logs)  # pylint: disable=protected-access
+    else:
+        return await sdk_async.get(request_id)
+
+
+@usage_lib.entrypoint
+async def terminate_replica(
+    service_name: str,
+    replica_id: int,
+    purge: bool,
+    stream_logs: Optional[
+        sdk_async.StreamConfig] = sdk_async.DEFAULT_STREAM_CONFIG
+) -> None:
+    """Async version of terminate_replica() that tears down a specific
+    replica."""
+    request_id = await context_utils.to_thread(sdk.terminate_replica,
+                                               service_name, replica_id, purge)
+    if stream_logs is not None:
+        return await sdk_async._stream_and_get(request_id, stream_logs)  # pylint: disable=protected-access
+    else:
+        return await sdk_async.get(request_id)
+
+
+@usage_lib.entrypoint
+async def status(
+    service_names: Optional[Union[str, List[str]]],
+    stream_logs: Optional[
+        sdk_async.StreamConfig] = sdk_async.DEFAULT_STREAM_CONFIG
+) -> List[Dict[str, Any]]:
+    """Async version of status() that sdk_async.gets service statuses."""
+    request_id = await context_utils.to_thread(sdk.status, service_names)
+    if stream_logs is not None:
+        return await sdk_async._stream_and_get(request_id, stream_logs)  # pylint: disable=protected-access
+    else:
+        return await sdk_async.get(request_id)
+
+
+@usage_lib.entrypoint
+async 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:
+    """Async version of tail_logs() that tails logs for a service."""
+    return await context_utils.to_thread(sdk.tail_logs, service_name, target,
+                                         replica_id, follow, output_stream)
+
+
+@usage_lib.entrypoint
+async def sync_down_logs(service_name: str,
+                         local_dir: str,
+                         *,
+                         targets: Optional[Union[
+                             str, 'serve_utils.ServiceComponent', List[Union[
+                                 str, 'serve_utils.ServiceComponent']]]] = None,
+                         replica_ids: Optional[List[int]] = None) -> None:
+    """Async version of sync_down_logs() that syncs down logs from service
+      components."""
+    return await context_utils.to_thread(sdk.sync_down_logs,
+                                         service_name,
+                                         local_dir,
+                                         targets=targets,
+                                         replica_ids=replica_ids)

sky/serve/constants.py

@@ -105,7 +105,8 @@ REPLICA_ID_ENV_VAR = 'SKYPILOT_SERVE_REPLICA_ID'
 # v1.0 - Introduce rolling update.
 # v2.0 - Added template-replica feature.
 # v3.0 - Added cluster pool.
-SERVE_VERSION = 3
+# v4.0 - Added pool argument to wait_service_registration.
+SERVE_VERSION = 4
 
 TERMINATE_REPLICA_VERSION_MISMATCH_ERROR = (
     'The version of service is outdated and does not support manually '

sky/serve/controller.py

@@ -4,6 +4,7 @@ Responsible for autoscaling and replica management.
 """
 import contextlib
 import logging
+import os
 import threading
 import time
 import traceback
@@ -242,7 +243,7 @@ class SkyServeController:
         threading.Thread(target=self._run_autoscaler).start()
 
         logger.info('SkyServe Controller started on '
-                    f'http://{self._host}:{self._port}')
+                    f'http://{self._host}:{self._port}. PID: {os.getpid()}')
 
         uvicorn.run(self._app, host=self._host, port=self._port)
 

sky/serve/load_balancer.py

@@ -1,6 +1,7 @@
 """LoadBalancer: Distribute any incoming request to all ready replicas."""
 import asyncio
 import logging
+import os
 import threading
 import traceback
 from typing import Dict, List, Optional, Union
@@ -254,7 +255,8 @@ class SkyServeLoadBalancer:
         protocol = 'https' if self._tls_credential is not None else 'http'
 
         logger.info('SkyServe Load Balancer started on '
-                    f'{protocol}://0.0.0.0:{self._load_balancer_port}')
+                    f'{protocol}://0.0.0.0:{self._load_balancer_port}. '
+                    f'PID: {os.getpid()}')
 
         uvicorn.run(self._app,
                     host='0.0.0.0',

sky/serve/serve_state.py

@@ -47,6 +47,10 @@ def create_table(cursor: 'sqlite3.Cursor', conn: 'sqlite3.Connection') -> None:
         service_name TEXT,
         spec BLOB,
         PRIMARY KEY (service_name, version))""")
+    cursor.execute("""\
+        CREATE TABLE IF NOT EXISTS ha_recovery_script (
+        service_name TEXT PRIMARY KEY,
+        script TEXT)""")
     conn.commit()
 
     # Backward compatibility.
@@ -71,6 +75,13 @@ def create_table(cursor: 'sqlite3.Cursor', conn: 'sqlite3.Connection') -> None:
     # Whether the service is a cluster pool.
     db_utils.add_column_to_table(cursor, conn, 'services', 'pool',
                                  'INTEGER DEFAULT 0')
+    # Add controller_pid for status tracking.
+    db_utils.add_column_to_table(cursor,
+                                 conn,
+                                 'services',
+                                 'controller_pid',
+                                 'INTEGER DEFAULT NULL',
+                                 value_to_replace_existing_entries=-1)
     conn.commit()
 
 
@@ -272,7 +283,8 @@ _SERVICE_STATUS_TO_COLOR = {
 @init_db
 def add_service(name: str, controller_job_id: int, policy: str,
                 requested_resources_str: str, load_balancing_policy: str,
-                status: ServiceStatus, tls_encrypted: bool, pool: bool) -> bool:
+                status: ServiceStatus, tls_encrypted: bool, pool: bool,
+                controller_pid: int) -> bool:
     """Add a service in the database.
 
     Returns:
@@ -287,11 +299,11 @@ def add_service(name: str, controller_job_id: int, policy: str,
                 INSERT INTO services
                 (name, controller_job_id, status, policy,
                 requested_resources_str, load_balancing_policy, tls_encrypted,
-                pool)
-                VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
+                pool, controller_pid)
+                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""",
                 (name, controller_job_id, status.value, policy,
                  requested_resources_str, load_balancing_policy,
-                 int(tls_encrypted), int(pool)))
+                 int(tls_encrypted), int(pool), controller_pid))
 
     except sqlite3.IntegrityError as e:
         if str(e) != _UNIQUE_CONSTRAINT_FAILED_ERROR_MSG:
@@ -300,6 +312,22 @@ def add_service(name: str, controller_job_id: int, policy: str,
     return True
 
 
+@init_db
+def update_service_controller_pid(service_name: str,
+                                  controller_pid: int) -> None:
+    """Updates the controller pid of a service.
+
+    This is used to update the controller pid of a service on ha recovery.
+    """
+    assert _DB_PATH is not None
+    with db_utils.safe_cursor(_DB_PATH) as cursor:
+        cursor.execute(
+            """\
+            UPDATE services SET
+            controller_pid=(?) WHERE name=(?)""",
+            (controller_pid, service_name))
+
+
 @init_db
 def remove_service(service_name: str) -> None:
     """Removes a service from the database."""
@@ -368,7 +396,8 @@ def set_service_load_balancer_port(service_name: str,
 def _get_service_from_row(row) -> Dict[str, Any]:
     (current_version, name, controller_job_id, controller_port,
      load_balancer_port, status, uptime, policy, _, _, requested_resources_str,
-     _, active_versions, load_balancing_policy, tls_encrypted, pool) = row[:16]
+     _, active_versions, load_balancing_policy, tls_encrypted, pool,
+     controller_pid) = row[:17]
     record = {
         'name': name,
         'controller_job_id': controller_job_id,
@@ -388,6 +417,7 @@ def _get_service_from_row(row) -> Dict[str, Any]:
         'load_balancing_policy': load_balancing_policy,
         'tls_encrypted': bool(tls_encrypted),
         'pool': bool(pool),
+        'controller_pid': controller_pid,
     }
     latest_spec = get_spec(name, current_version)
     if latest_spec is not None:
@@ -666,3 +696,38 @@ def get_service_load_balancer_port(service_name: str) -> int:
         if row is None:
             raise ValueError(f'Service {service_name} does not exist.')
         return row[0]
+
+
+@init_db
+def get_ha_recovery_script(service_name: str) -> Optional[str]:
+    """Gets the HA recovery script for a service."""
+    assert _DB_PATH is not None
+    with db_utils.safe_cursor(_DB_PATH) as cursor:
+        cursor.execute(
+            'SELECT script FROM ha_recovery_script WHERE service_name = ?',
+            (service_name,))
+        row = cursor.fetchone()
+        if row is None:
+            return None
+        return row[0]
+
+
+@init_db
+def set_ha_recovery_script(service_name: str, script: str) -> None:
+    """Sets the HA recovery script for a service."""
+    assert _DB_PATH is not None
+    with db_utils.safe_cursor(_DB_PATH) as cursor:
+        cursor.execute(
+            """\
+            INSERT OR REPLACE INTO ha_recovery_script
+            (service_name, script)
+            VALUES (?, ?)""", (service_name, script))
+
+
+@init_db
+def remove_ha_recovery_script(service_name: str) -> None:
+    """Removes the HA recovery script for a service."""
+    assert _DB_PATH is not None
+    with db_utils.safe_cursor(_DB_PATH) as cursor:
+        cursor.execute('DELETE FROM ha_recovery_script WHERE service_name = ?',
+                       (service_name,))

sky/serve/serve_utils.py

@@ -2,6 +2,7 @@
 import base64
 import collections
 import dataclasses
+import datetime
 import enum
 import os
 import pathlib
@@ -33,6 +34,7 @@ from sky.serve import spot_placer
 from sky.skylet import constants as skylet_constants
 from sky.skylet import job_lib
 from sky.utils import annotations
+from sky.utils import command_runner
 from sky.utils import common_utils
 from sky.utils import log_utils
 from sky.utils import message_utils
@@ -258,13 +260,76 @@ def get_service_filelock_path(pool: str) -> str:
 
 
 @annotations.lru_cache(scope='request', maxsize=1)
-def is_consolidation_mode() -> bool:
+def is_consolidation_mode(pool: bool = False) -> bool:
+    # Use jobs config for pool consolidation mode.
+    controller_type = 'jobs' if pool else 'serve'
     consolidation_mode = skypilot_config.get_nested(
-        ('serve', 'controller', 'consolidation_mode'), default_value=False)
-    # _check_consolidation_mode_consistency(consolidation_mode)
+        (controller_type, 'controller', 'consolidation_mode'),
+        default_value=False)
+    # _check_consolidation_mode_consistency(consolidation_mode, pool)
     return consolidation_mode
 
 
+def ha_recovery_for_consolidation_mode(pool: bool):
+    """Recovery logic for HA mode."""
+    # No setup recovery is needed in consolidation mode, as the API server
+    # already has all runtime installed. Directly start jobs recovery here.
+    # Refers to sky/templates/kubernetes-ray.yml.j2 for more details.
+    runner = command_runner.LocalProcessCommandRunner()
+    noun = 'pool' if pool else 'serve'
+    capnoun = noun.capitalize()
+    prefix = f'{noun}_'
+    with open(skylet_constants.HA_PERSISTENT_RECOVERY_LOG_PATH.format(prefix),
+              'w',
+              encoding='utf-8') as f:
+        start = time.time()
+        f.write(f'Starting HA recovery at {datetime.datetime.now()}\n')
+        for service_name in serve_state.get_glob_service_names(None):
+            svc = _get_service_status(service_name,
+                                      pool=pool,
+                                      with_replica_info=False)
+            if svc is None:
+                continue
+            controller_pid = svc['controller_pid']
+            if controller_pid is not None:
+                try:
+                    if _controller_process_alive(controller_pid, service_name):
+                        f.write(f'Controller pid {controller_pid} for '
+                                f'{noun} {service_name} is still running. '
+                                'Skipping recovery.\n')
+                        continue
+                except Exception:  # pylint: disable=broad-except
+                    # _controller_process_alive may raise if psutil fails; we
+                    # should not crash the recovery logic because of this.
+                    f.write('Error checking controller pid '
+                            f'{controller_pid} for {noun} {service_name}\n')
+
+            script = serve_state.get_ha_recovery_script(service_name)
+            if script is None:
+                f.write(f'{capnoun} {service_name}\'s recovery script does '
+                        'not exist. Skipping recovery.\n')
+                continue
+            rc, out, err = runner.run(script, require_outputs=True)
+            if rc:
+                f.write(f'Recovery script returned {rc}. '
+                        f'Output: {out}\nError: {err}\n')
+            f.write(f'{capnoun} {service_name} completed recovery at '
+                    f'{datetime.datetime.now()}\n')
+        f.write(f'HA recovery completed at {datetime.datetime.now()}\n')
+        f.write(f'Total recovery time: {time.time() - start} seconds\n')
+
+
+def _controller_process_alive(pid: int, service_name: str) -> bool:
+    """Check if the controller process is alive."""
+    try:
+        process = psutil.Process(pid)
+        cmd_str = ' '.join(process.cmdline())
+        return process.is_running(
+        ) and f'--service-name {service_name}' in cmd_str
+    except psutil.NoSuchProcess:
+        return False
+
+
 def validate_service_task(task: 'sky.Task', pool: bool) -> None:
     """Validate the task for Sky Serve.
 
@@ -460,22 +525,53 @@ def set_service_status_and_active_versions_from_replica(
         active_versions=active_versions)
 
 
-def update_service_status() -> None:
-    if is_consolidation_mode():
-        # TODO(tian): PID-based tracking.
-        return
-    services = serve_state.get_services()
-    for record in services:
-        if record['status'] == serve_state.ServiceStatus.SHUTTING_DOWN:
+def update_service_status(pool: bool) -> None:
+    noun = 'pool' if pool else 'serve'
+    capnoun = noun.capitalize()
+    service_names = serve_state.get_glob_service_names(None)
+    for service_name in service_names:
+        record = _get_service_status(service_name,
+                                     pool=pool,
+                                     with_replica_info=False)
+        if record is None:
+            continue
+        service_status = record['status']
+        if service_status == serve_state.ServiceStatus.SHUTTING_DOWN:
             # Skip services that is shutting down.
             continue
-        controller_job_id = record['controller_job_id']
-        assert controller_job_id is not None
-        controller_status = job_lib.get_status(controller_job_id)
-        if controller_status is None or controller_status.is_terminal():
-            # If controller job is not running, set it as controller failed.
-            serve_state.set_service_status_and_active_versions(
-                record['name'], serve_state.ServiceStatus.CONTROLLER_FAILED)
+
+        logger.info(f'Update {noun} status for {service_name!r} '
+                    f'with status {service_status}')
+
+        controller_pid = record['controller_pid']
+        if controller_pid is None:
+            logger.info(f'{capnoun} {service_name!r} controller pid is None. '
+                        f'Unexpected status {service_status}. Set to failure.')
+        elif controller_pid < 0:
+            # Backwards compatibility: this service was submitted when ray was
+            # still used for controller process management. We set the
+            # value_to_replace_existing_entries to -1 to indicate historical
+            # services.
+            # TODO(tian): Remove before 0.13.0.
+            controller_job_id = record['controller_job_id']
+            assert controller_job_id is not None
+            controller_status = job_lib.get_status(controller_job_id)
+            if (controller_status is not None and
+                    not controller_status.is_terminal()):
+                continue
+            logger.info(f'Updating {noun} {service_name!r} in old version. '
+                        f'SkyPilot job status: {controller_status}. '
+                        'Set to failure.')
+        else:
+            if _controller_process_alive(controller_pid, service_name):
+                # The controller is still running.
+                continue
+            logger.info(f'{capnoun} {service_name!r} controller pid '
+                        f'{controller_pid} is not alive. Set to failure.')
+
+        # If controller job is not running, set it as controller failed.
+        serve_state.set_service_status_and_active_versions(
+            service_name, serve_state.ServiceStatus.CONTROLLER_FAILED)
 
 
 def update_service_encoded(service_name: str, version: int, mode: str,
@@ -754,9 +850,11 @@ def _terminate_failed_services(
     shutil.rmtree(service_dir)
     serve_state.remove_service(service_name)
     serve_state.delete_all_versions(service_name)
+    serve_state.remove_ha_recovery_script(service_name)
 
     if not remaining_replica_clusters:
         return None
+    # TODO(tian): Try to terminate those replica clusters.
     remaining_identity = ', '.join(remaining_replica_clusters)
     return (f'{colorama.Fore.YELLOW}terminate service {service_name!r} with '
             f'failed status ({service_status}). This may indicate a resource '
@@ -845,7 +943,8 @@ def terminate_services(service_names: Optional[List[str]], purge: bool,
     return '\n'.join(messages)
 
 
-def wait_service_registration(service_name: str, job_id: int) -> str:
+def wait_service_registration(service_name: str, job_id: int,
+                              pool: bool) -> str:
     """Util function to call at the end of `sky.serve.up()`.
 
     This function will:
@@ -862,7 +961,7 @@ def wait_service_registration(service_name: str, job_id: int) -> str:
     setup_completed = False
     while True:
         # TODO(tian): PID-based tracking.
-        if not is_consolidation_mode():
+        if not is_consolidation_mode(pool):
             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
@@ -888,7 +987,7 @@ def wait_service_registration(service_name: str, job_id: int) -> str:
         record = serve_state.get_service_from_name(service_name)
         if record is not None:
             # TODO(tian): PID-based tracking.
-            if (not is_consolidation_mode() and
+            if (not is_consolidation_mode(pool) and
                     job_id != record['controller_job_id']):
                 with ux_utils.print_exception_no_traceback():
                     raise ValueError(
@@ -1420,10 +1519,13 @@ class ServeCodeGen:
         return cls._build(code)
 
     @classmethod
-    def wait_service_registration(cls, service_name: str, job_id: int) -> str:
+    def wait_service_registration(cls, service_name: str, job_id: int,
+                                  pool: bool) -> str:
         code = [
+            f'kwargs={{}} if serve_version < 4 else {{"pool": {pool}}}',
             'msg = serve_utils.wait_service_registration('
-            f'{service_name!r}, {job_id})', 'print(msg, end="", flush=True)'
+            f'{service_name!r}, {job_id}, **kwargs)',
+            'print(msg, end="", flush=True)'
         ]
         return cls._build(code)
 

sky/serve/server/impl.py

@@ -102,10 +102,10 @@ def up(
     pool: bool = False,
 ) -> Tuple[str, str]:
     """Spins up a service or a pool."""
-    if pool and not serve_utils.is_consolidation_mode():
+    if pool and not serve_utils.is_consolidation_mode(pool):
         raise ValueError(
             'Pool is only supported in consolidation mode. To fix, set '
-            '`serve.controller.consolidation_mode: true` in SkyPilot config.')
+            '`jobs.controller.consolidation_mode: true` in SkyPilot config.')
     task.validate()
     serve_utils.validate_service_task(task, pool=pool)
     assert task.service is not None
@@ -174,7 +174,8 @@ def up(
             prefix=f'controller-task-{service_name}-',
             mode='w',
     ) as controller_file:
-        controller_name = common.SKY_SERVE_CONTROLLER_NAME
+        controller = controller_utils.get_controller_for_pool(pool)
+        controller_name = controller.value.cluster_name
         task_config = task.to_yaml_config()
         common_utils.dump_yaml(service_file.name, task_config)
         remote_tmp_task_yaml_path = (
@@ -187,7 +188,7 @@ def up(
             controller=controller_utils.Controllers.SKY_SERVE_CONTROLLER,
             task_resources=task.resources)
         controller_job_id = None
-        if serve_utils.is_consolidation_mode():
+        if serve_utils.is_consolidation_mode(pool):
             controller_job_id = 0
 
         vars_to_fill = {
@@ -238,7 +239,7 @@ def up(
         # 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():
+        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():
@@ -251,9 +252,9 @@ def up(
                         _disable_controller_check=True,
                     )
         else:
+            controller_type = controller_utils.get_controller_for_pool(pool)
             controller_handle = backend_utils.is_controller_accessible(
-                controller=controller_utils.Controllers.SKY_SERVE_CONTROLLER,
-                stopped_message='')
+                controller=controller_type, stopped_message='')
             backend = backend_utils.get_backend_from_handle(controller_handle)
             assert isinstance(backend, backends.CloudVmRayBackend)
             backend.sync_file_mounts(
@@ -270,10 +271,8 @@ def up(
             ]
             run_script = '\n'.join(env_cmds + [run_script])
             # Dump script for high availability recovery.
-            # if controller_utils.high_availability_specified(
-            #         controller_name):
-            #     managed_job_state.set_ha_recovery_script(
-            #         consolidation_mode_job_id, run_script)
+            if controller_utils.high_availability_specified(controller_name):
+                serve_state.set_ha_recovery_script(service_name, run_script)
             backend.run_on_head(controller_handle, run_script)
 
         style = colorama.Style
@@ -289,7 +288,7 @@ def up(
             # and return the endpoint if the job id matches. Otherwise it will
             # return None.
             code = serve_utils.ServeCodeGen.wait_service_registration(
-                service_name, controller_job_id)
+                service_name, controller_job_id, pool)
             backend = backend_utils.get_backend_from_handle(controller_handle)
             assert isinstance(backend, backends.CloudVmRayBackend)
             assert isinstance(controller_handle,
@@ -304,7 +303,7 @@ def up(
                 returncode, code, f'Failed to wait for {noun} initialization',
                 lb_port_payload)
         except exceptions.CommandError:
-            if serve_utils.is_consolidation_mode():
+            if serve_utils.is_consolidation_mode(pool):
                 with ux_utils.print_exception_no_traceback():
                     raise RuntimeError(
                         f'Failed to wait for {noun} initialization. '
@@ -339,7 +338,7 @@ def up(
         else:
             lb_port = serve_utils.load_service_initialization_result(
                 lb_port_payload)
-            if not serve_utils.is_consolidation_mode():
+            if not serve_utils.is_consolidation_mode(pool):
                 socket_endpoint = backend_utils.get_endpoints(
                     controller_handle.cluster_name,
                     lb_port,
@@ -442,8 +441,9 @@ def update(
                        'effect. To update TLS keyfile and certfile, please '
                        'tear down the service and spin up a new one.')
 
+    controller_type = controller_utils.get_controller_for_pool(pool)
     handle = backend_utils.is_controller_accessible(
-        controller=controller_utils.Controllers.SKY_SERVE_CONTROLLER,
+        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}'
@@ -572,9 +572,9 @@ def apply(
     """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_utils.Controllers.SKY_SERVE_CONTROLLER,
-                stopped_message='')
+                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,
@@ -598,8 +598,9 @@ def down(
         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_utils.Controllers.SKY_SERVE_CONTROLLER,
+        controller=controller_type,
         stopped_message=f'All {noun}s should have terminated.')
 
     service_names_str = ','.join(service_names)
@@ -624,7 +625,7 @@ def down(
     except exceptions.FetchClusterInfoError as e:
         raise RuntimeError(
             'Failed to fetch controller IP. Please refresh controller status '
-            f'by `sky status -r {common.SKY_SERVE_CONTROLLER_NAME}` '
+            f'by `sky status -r {controller_type.value.cluster_name}` '
             'and try again.') from e
 
     try:
@@ -654,7 +655,7 @@ def status(
             raise RuntimeError(f'Failed to refresh {noun}s status '
                                'due to network error.') from e
 
-    controller_type = controller_utils.Controllers.SKY_SERVE_CONTROLLER
+    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.
@@ -690,7 +691,7 @@ def status(
         if service_record['load_balancer_port'] is not None:
             try:
                 lb_port = service_record['load_balancer_port']
-                if not serve_utils.is_consolidation_mode():
+                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)