skypilot-nightly 1.0.0.dev20250502__py3-none-any.whl → 1.0.0.dev20251203__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

@@ -65,7 +65,8 @@ AUTOSCALER_DEFAULT_DOWNSCALE_DELAY_SECONDS = 1200
 # TODO(tian): We might need to be careful that service logs can take a lot of
 # disk space. Maybe we could use a larger disk size, migrate to cloud storage or
 # do some log rotation.
-CONTROLLER_RESOURCES = {'cpus': '4+', 'disk_size': 200}
+# Set default minimal memory to 8GB to allow at least one service to run.
+CONTROLLER_RESOURCES = {'cpus': '4+', 'memory': '8+', 'disk_size': 200}
 # Autostop config for the jobs controller. These are the default values for
 # serve.controller.autostop in ~/.sky/config.yaml.
 CONTROLLER_AUTOSTOP = {
@@ -73,13 +74,6 @@ CONTROLLER_AUTOSTOP = {
     'down': False,
 }
 
-# Due to the CPU/memory usage of the controller process launched with a job on
-# controller VM (use ray job under the hood), we need to reserve some CPU/memory
-# for each serve controller process.
-# Serve: A default controller with 4 vCPU and 16 GB memory can run up to 16
-# services.
-CONTROLLER_MEMORY_USAGE_GB = 1.0
-
 # A period of time to initialize your service. Any readiness probe failures
 # during this period will be ignored.
 DEFAULT_INITIAL_DELAY_SECONDS = 1200
@@ -104,8 +98,17 @@ REPLICA_ID_ENV_VAR = 'SKYPILOT_SERVE_REPLICA_ID'
 # Changelog:
 # v1.0 - Introduce rolling update.
 # v2.0 - Added template-replica feature.
-SERVE_VERSION = 2
+# v3.0 - Added pool.
+# v4.0 - Added pool argument to wait_service_registration.
+# v5.0 - Added pool argument to stream_serve_process_logs & stream_replica_logs.
+SERVE_VERSION = 5
 
 TERMINATE_REPLICA_VERSION_MISMATCH_ERROR = (
     'The version of service is outdated and does not support manually '
     'terminating replicas. Please terminate the service and spin up again.')
+
+# Dummy run command for pool.
+POOL_DUMMY_RUN_COMMAND = 'echo "setup done"'
+
+# Error message for max number of services reached.
+MAX_NUMBER_OF_SERVICES_REACHED_ERROR = 'Max number of services reached.'

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
@@ -20,17 +21,20 @@ from sky.serve import autoscalers
 from sky.serve import replica_managers
 from sky.serve import serve_state
 from sky.serve import serve_utils
+from sky.skylet import constants
 from sky.utils import common_utils
+from sky.utils import context_utils
 from sky.utils import ux_utils
 
 logger = sky_logging.init_logger(__name__)
 
 
-class SuppressSuccessGetAccessLogsFilter(logging.Filter):
+class AutoscalerInfoFilter(logging.Filter):
 
     def filter(self, record: logging.LogRecord) -> bool:
         message = record.getMessage()
-        return not ('GET' in message and '200' in message)
+        return not ('GET' in message and '200' in message and
+                    '/autoscaler/info' in message)
 
 
 class SkyServeController:
@@ -42,12 +46,12 @@ class SkyServeController:
     """
 
     def __init__(self, service_name: str, service_spec: serve.SkyServiceSpec,
-                 task_yaml: str, host: str, port: int) -> None:
+                 version: int, host: str, port: int) -> None:
         self._service_name = service_name
         self._replica_manager: replica_managers.ReplicaManager = (
             replica_managers.SkyPilotReplicaManager(service_name=service_name,
                                                     spec=service_spec,
-                                                    task_yaml_path=task_yaml))
+                                                    version=version))
         self._autoscaler: autoscalers.Autoscaler = (
             autoscalers.Autoscaler.from_spec(service_name, service_spec))
         self._host = host
@@ -59,6 +63,7 @@ class SkyServeController:
         uvicorn_access_logger = logging.getLogger('uvicorn.access')
         for handler in uvicorn_access_logger.handlers:
             handler.setFormatter(sky_logging.FORMATTER)
+            handler.addFilter(AutoscalerInfoFilter())
         yield
 
     def _run_autoscaler(self):
@@ -74,7 +79,11 @@ class SkyServeController:
                 assert record is not None, ('No service record found for '
                                             f'{self._service_name}')
                 active_versions = record['active_versions']
-                logger.info(f'All replica info: {replica_infos}')
+                logger.info(f'All replica info for autoscaler: {replica_infos}')
+
+                # Autoscaler now extracts GPU type info directly from
+                # replica_infos in generate_scaling_decisions method
+                # for better decoupling.
                 scaling_options = self._autoscaler.generate_scaling_decisions(
                     replica_infos, active_versions)
                 for scaling_option in scaling_options:
@@ -99,6 +108,11 @@ class SkyServeController:
 
     def run(self) -> None:
 
+        @self._app.get('/autoscaler/info')
+        async def get_autoscaler_info() -> fastapi.Response:
+            return responses.JSONResponse(content=self._autoscaler.info(),
+                                          status_code=200)
+
         @self._app.post('/controller/load_balancer_sync')
         async def load_balancer_sync(
                 request: fastapi.Request) -> fastapi.Response:
@@ -109,11 +123,37 @@ class SkyServeController:
             timestamps: List[int] = request_aggregator.get('timestamps', [])
             logger.info(f'Received {len(timestamps)} inflight requests.')
             self._autoscaler.collect_request_information(request_aggregator)
-            return responses.JSONResponse(content={
-                'ready_replica_urls':
-                    self._replica_manager.get_active_replica_urls()
-            },
-                                          status_code=200)
+
+            # Get replica information for instance-aware load balancing
+            replica_infos = serve_state.get_replica_infos(self._service_name)
+            ready_replica_urls = self._replica_manager.get_active_replica_urls()
+
+            # Use URL-to-info mapping to avoid duplication
+            replica_info = {}
+            for info in replica_infos:
+                if info.url in ready_replica_urls:
+                    # Get GPU type from handle.launched_resources.accelerators
+                    gpu_type = 'unknown'
+                    handle = info.handle()
+                    if handle is not None:
+                        accelerators = handle.launched_resources.accelerators
+                        if accelerators and len(accelerators) > 0:
+                            # Get the first accelerator type
+                            gpu_type = list(accelerators.keys())[0]
+
+                    replica_info[info.url] = {'gpu_type': gpu_type}
+
+            # Check that all ready replica URLs are included in replica_info
+            missing_urls = set(ready_replica_urls) - set(replica_info.keys())
+            if missing_urls:
+                logger.warning(f'Ready replica URLs missing from replica_info: '
+                               f'{missing_urls}')
+                # fallback: add missing URLs with unknown GPU type
+                for url in missing_urls:
+                    replica_info[url] = {'gpu_type': 'unknown'}
+
+            return responses.JSONResponse(
+                content={'replica_info': replica_info}, status_code=200)
 
         @self._app.post('/controller/update_service')
         async def update_service(request: fastapi.Request) -> fastapi.Response:
@@ -133,7 +173,11 @@ class SkyServeController:
                 # See sky/serve/core.py::update
                 latest_task_yaml = serve_utils.generate_task_yaml_file_name(
                     self._service_name, version)
-                service = serve.SkyServiceSpec.from_yaml(latest_task_yaml)
+                with open(latest_task_yaml, 'r', encoding='utf-8') as f:
+                    yaml_content = f.read()
+                service = serve.SkyServiceSpec.from_yaml_str(yaml_content)
+                serve_state.add_or_update_version(self._service_name, version,
+                                                  service, yaml_content)
                 logger.info(
                     f'Update to new version version {version}: {service}')
 
@@ -155,9 +199,13 @@ class SkyServeController:
                 return responses.JSONResponse(content={'message': 'Success'},
                                               status_code=200)
             except Exception as e:  # pylint: disable=broad-except
-                logger.error(f'Error in update_service: '
-                             f'{common_utils.format_exception(e)}')
-                return responses.JSONResponse(content={'message': 'Error'},
+                exception_str = common_utils.format_exception(e)
+                logger.error(f'Error in update_service: {exception_str}')
+                return responses.JSONResponse(content={
+                    'message': 'Error',
+                    'exception': exception_str,
+                    'traceback': traceback.format_exc()
+                },
                                               status_code=500)
 
         @self._app.post('/controller/terminate_replica')
@@ -232,7 +280,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)
 
@@ -240,7 +288,10 @@ class SkyServeController:
 # TODO(tian): Probably we should support service that will stop the VM in
 # specific time period.
 def run_controller(service_name: str, service_spec: serve.SkyServiceSpec,
-                   task_yaml: str, controller_host: str, controller_port: int):
-    controller = SkyServeController(service_name, service_spec, task_yaml,
+                   version: int, controller_host: str, controller_port: int):
+    os.environ[constants.OVERRIDE_CONSOLIDATION_MODE] = 'true'
+    # Hijack sys.stdout/stderr to be context aware.
+    context_utils.hijack_sys_attrs()
+    controller = SkyServeController(service_name, service_spec, version,
                                     controller_host, controller_port)
     controller.run()

sky/serve/load_balancer.py

@@ -1,8 +1,10 @@
 """LoadBalancer: Distribute any incoming request to all ready replicas."""
 import asyncio
 import logging
+import os
 import threading
-from typing import Dict, Optional, Union
+import traceback
+from typing import Dict, List, Optional, Union
 
 import aiohttp
 import fastapi
@@ -28,11 +30,13 @@ class SkyServeLoadBalancer:
     """
 
     def __init__(
-            self,
-            controller_url: str,
-            load_balancer_port: int,
-            load_balancing_policy_name: Optional[str] = None,
-            tls_credential: Optional[serve_utils.TLSCredential] = None) -> None:
+        self,
+        controller_url: str,
+        load_balancer_port: int,
+        load_balancing_policy_name: Optional[str] = None,
+        tls_credential: Optional[serve_utils.TLSCredential] = None,
+        target_qps_per_replica: Optional[Union[float, Dict[str, float]]] = None
+    ) -> None:
         """Initialize the load balancer.
 
         Args:
@@ -42,6 +46,9 @@ class SkyServeLoadBalancer:
                 to use. Defaults to None.
             tls_credentials: The TLS credentials for HTTPS endpoint. Defaults
                 to None.
+            target_qps_per_replica: Target QPS per replica for instance-aware
+                load balancing. Can be a float or dict mapping GPU types to QPS.
+                Defaults to None.
         """
         self._app = fastapi.FastAPI()
         self._controller_url: str = controller_url
@@ -49,6 +56,15 @@ class SkyServeLoadBalancer:
         # Use the registry to create the load balancing policy
         self._load_balancing_policy = lb_policies.LoadBalancingPolicy.make(
             load_balancing_policy_name)
+
+        # Set accelerator QPS for instance-aware policies
+        if (target_qps_per_replica and
+                isinstance(target_qps_per_replica, dict) and
+                isinstance(self._load_balancing_policy,
+                           lb_policies.InstanceAwareLeastLoadPolicy)):
+            self._load_balancing_policy.set_target_qps_per_accelerator(
+                target_qps_per_replica)
+
         logger.info('Starting load balancer with policy '
                     f'{load_balancing_policy_name}.')
         self._request_aggregator: serve_utils.RequestsAggregator = (
@@ -69,6 +85,56 @@ class SkyServeLoadBalancer:
         # updating it from _sync_with_controller.
         self._client_pool_lock: threading.Lock = threading.Lock()
 
+    async def _sync_with_controller_once(self) -> List[asyncio.Task]:
+        close_client_tasks = []
+        ready_replica_urls = []
+        replica_info = {}
+
+        async with aiohttp.ClientSession() as session:
+            try:
+                # Send request information
+                async with session.post(
+                        self._controller_url + '/controller/load_balancer_sync',
+                        json={
+                            'request_aggregator':
+                                self._request_aggregator.to_dict()
+                        },
+                        timeout=aiohttp.ClientTimeout(5),
+                ) as response:
+                    # Clean up after reporting request info to avoid OOM.
+                    self._request_aggregator.clear()
+                    response.raise_for_status()
+                    response_json = await response.json()
+                    replica_info = response_json.get('replica_info', {})
+                    ready_replica_urls = list(replica_info.keys())
+            except (aiohttp.ClientError, asyncio.TimeoutError) as e:
+                logger.error(f'An error occurred when syncing with '
+                             f'the controller: {e}'
+                             f'\nTraceback: {traceback.format_exc()}')
+            else:
+                logger.info(f'Available Replica URLs: {ready_replica_urls}')
+                with self._client_pool_lock:
+                    self._load_balancing_policy.set_ready_replicas(
+                        ready_replica_urls)
+                    # Set replica info for instance-aware policies
+                    if isinstance(self._load_balancing_policy,
+                                  lb_policies.InstanceAwareLeastLoadPolicy):
+                        self._load_balancing_policy.set_replica_info(
+                            replica_info)
+                    for replica_url in ready_replica_urls:
+                        if replica_url not in self._client_pool:
+                            self._client_pool[replica_url] = httpx.AsyncClient(
+                                base_url=replica_url)
+                    urls_to_close = set(
+                        self._client_pool.keys()) - set(ready_replica_urls)
+                    client_to_close = []
+                    for replica_url in urls_to_close:
+                        client_to_close.append(
+                            self._client_pool.pop(replica_url))
+                for client in client_to_close:
+                    close_client_tasks.append(client.aclose())
+        return close_client_tasks
+
     async def _sync_with_controller(self):
         """Sync with controller periodically.
 
@@ -82,49 +148,16 @@ class SkyServeLoadBalancer:
         await asyncio.sleep(5)
 
         while True:
-            close_client_tasks = []
-            async with aiohttp.ClientSession() as session:
-                try:
-                    # Send request information
-                    async with session.post(
-                            self._controller_url +
-                            '/controller/load_balancer_sync',
-                            json={
-                                'request_aggregator':
-                                    self._request_aggregator.to_dict()
-                            },
-                            timeout=aiohttp.ClientTimeout(5),
-                    ) as response:
-                        # Clean up after reporting request info to avoid OOM.
-                        self._request_aggregator.clear()
-                        response.raise_for_status()
-                        response_json = await response.json()
-                        ready_replica_urls = response_json.get(
-                            'ready_replica_urls', [])
-                except aiohttp.ClientError as e:
-                    logger.error('An error occurred when syncing with '
-                                 f'the controller: {e}')
-                else:
-                    logger.info(f'Available Replica URLs: {ready_replica_urls}')
-                    with self._client_pool_lock:
-                        self._load_balancing_policy.set_ready_replicas(
-                            ready_replica_urls)
-                        for replica_url in ready_replica_urls:
-                            if replica_url not in self._client_pool:
-                                self._client_pool[replica_url] = (
-                                    httpx.AsyncClient(base_url=replica_url))
-                        urls_to_close = set(
-                            self._client_pool.keys()) - set(ready_replica_urls)
-                        client_to_close = []
-                        for replica_url in urls_to_close:
-                            client_to_close.append(
-                                self._client_pool.pop(replica_url))
-                    for client in client_to_close:
-                        close_client_tasks.append(client.aclose())
-
-            await asyncio.sleep(constants.LB_CONTROLLER_SYNC_INTERVAL_SECONDS)
-            # Await those tasks after the interval to avoid blocking.
-            await asyncio.gather(*close_client_tasks)
+            try:
+                close_client_tasks = await self._sync_with_controller_once()
+                await asyncio.sleep(
+                    constants.LB_CONTROLLER_SYNC_INTERVAL_SECONDS)
+                # Await those tasks after the interval to avoid blocking.
+                await asyncio.gather(*close_client_tasks)
+            except Exception as e:  # pylint: disable=broad-except
+                logger.error(f'An error occurred when syncing with '
+                             f'the controller: {e}'
+                             f'\nTraceback: {traceback.format_exc()}')
 
     async def _proxy_request_to(
         self, url: str, request: fastapi.Request
@@ -168,7 +201,8 @@ class SkyServeLoadBalancer:
                 background=background.BackgroundTask(background_func))
         except (httpx.RequestError, httpx.HTTPStatusError) as e:
             logger.error(f'Error when proxy request to {url}: '
-                         f'{common_utils.format_exception(e)}')
+                         f'{common_utils.format_exception(e)}'
+                         f'\nTraceback: {traceback.format_exc()}')
             return e
 
     async def _proxy_with_retries(
@@ -243,7 +277,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',
@@ -252,23 +287,31 @@ class SkyServeLoadBalancer:
 
 
 def run_load_balancer(
-        controller_addr: str,
-        load_balancer_port: int,
-        load_balancing_policy_name: Optional[str] = None,
-        tls_credential: Optional[serve_utils.TLSCredential] = None) -> None:
+    controller_addr: str,
+    load_balancer_port: int,
+    load_balancing_policy_name: Optional[str] = None,
+    tls_credential: Optional[serve_utils.TLSCredential] = None,
+    target_qps_per_replica: Optional[Union[float, Dict[str, float]]] = None
+) -> None:
     """ Run the load balancer.
 
     Args:
         controller_addr: The address of the controller.
         load_balancer_port: The port where the load balancer listens to.
-        policy_name: The name of the load balancing policy to use. Defaults to
-            None.
+        policy_name: The name of the load balancing policy to use.
+        Defaults to None.
+        tls_credential:
+            The TLS credentials for HTTPS endpoint. Defaults to None.
+        target_qps_per_replica: Target QPS per replica for instance-aware
+            load balancing. Can be a float or dict mapping GPU types to QPS.
+            Defaults to None.
     """
     load_balancer = SkyServeLoadBalancer(
         controller_url=controller_addr,
         load_balancer_port=load_balancer_port,
         load_balancing_policy_name=load_balancing_policy_name,
-        tls_credential=tls_credential)
+        tls_credential=tls_credential,
+        target_qps_per_replica=target_qps_per_replica)
     load_balancer.run()
 
 
@@ -292,5 +335,8 @@ if __name__ == '__main__':
         help=f'The load balancing policy to use. Available policies: '
         f'{", ".join(available_policies)}.')
     args = parser.parse_args()
-    run_load_balancer(args.controller_addr, args.load_balancer_port,
-                      args.load_balancing_policy)
+    run_load_balancer(args.controller_addr,
+                      args.load_balancer_port,
+                      args.load_balancing_policy,
+                      tls_credential=None,
+                      target_qps_per_replica=None)