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

sky/serve/server/server.py

@@ -0,0 +1,112 @@
+"""Rest APIs for SkyServe."""
+
+import fastapi
+
+from sky import sky_logging
+from sky.serve.server import core
+from sky.server import stream_utils
+from sky.server.requests import executor
+from sky.server.requests import payloads
+from sky.server.requests import requests as api_requests
+from sky.utils import common
+
+logger = sky_logging.init_logger(__name__)
+router = fastapi.APIRouter()
+
+
+@router.post('/up')
+async def up(
+    request: fastapi.Request,
+    up_body: payloads.ServeUpBody,
+) -> None:
+    executor.schedule_request(
+        request_id=request.state.request_id,
+        request_name='serve.up',
+        request_body=up_body,
+        func=core.up,
+        schedule_type=api_requests.ScheduleType.LONG,
+        request_cluster_name=common.SKY_SERVE_CONTROLLER_NAME,
+    )
+
+
+@router.post('/update')
+async def update(
+    request: fastapi.Request,
+    update_body: payloads.ServeUpdateBody,
+) -> None:
+    executor.schedule_request(
+        request_id=request.state.request_id,
+        request_name='serve.update',
+        request_body=update_body,
+        func=core.update,
+        schedule_type=api_requests.ScheduleType.SHORT,
+        request_cluster_name=common.SKY_SERVE_CONTROLLER_NAME,
+    )
+
+
+@router.post('/down')
+async def down(
+    request: fastapi.Request,
+    down_body: payloads.ServeDownBody,
+) -> None:
+    executor.schedule_request(
+        request_id=request.state.request_id,
+        request_name='serve.down',
+        request_body=down_body,
+        func=core.down,
+        schedule_type=api_requests.ScheduleType.SHORT,
+        request_cluster_name=common.SKY_SERVE_CONTROLLER_NAME,
+    )
+
+
+@router.post('/terminate-replica')
+async def terminate_replica(
+    request: fastapi.Request,
+    terminate_replica_body: payloads.ServeTerminateReplicaBody,
+) -> None:
+    executor.schedule_request(
+        request_id=request.state.request_id,
+        request_name='serve.terminate_replica',
+        request_body=terminate_replica_body,
+        func=core.terminate_replica,
+        schedule_type=api_requests.ScheduleType.SHORT,
+        request_cluster_name=common.SKY_SERVE_CONTROLLER_NAME,
+    )
+
+
+@router.post('/status')
+async def status(
+    request: fastapi.Request,
+    status_body: payloads.ServeStatusBody,
+) -> None:
+    executor.schedule_request(
+        request_id=request.state.request_id,
+        request_name='serve.status',
+        request_body=status_body,
+        func=core.status,
+        schedule_type=api_requests.ScheduleType.SHORT,
+        request_cluster_name=common.SKY_SERVE_CONTROLLER_NAME,
+    )
+
+
+@router.post('/logs')
+async def tail_logs(
+    request: fastapi.Request, log_body: payloads.ServeLogsBody,
+    background_tasks: fastapi.BackgroundTasks
+) -> fastapi.responses.StreamingResponse:
+    executor.schedule_request(
+        request_id=request.state.request_id,
+        request_name='serve.logs',
+        request_body=log_body,
+        func=core.tail_logs,
+        schedule_type=api_requests.ScheduleType.SHORT,
+        request_cluster_name=common.SKY_SERVE_CONTROLLER_NAME,
+    )
+
+    request_task = api_requests.get_request(request.state.request_id)
+
+    return stream_utils.stream_response(
+        request_id=request_task.request_id,
+        logs_path=request_task.log_path,
+        background_tasks=background_tasks,
+    )

sky/serve/service.py

@@ -9,7 +9,7 @@ import pathlib
 import shutil
 import time
 import traceback
-from typing import Dict, List
+from typing import Dict
 
 import filelock
 
@@ -73,6 +73,12 @@ def cleanup_storage(task_yaml: str) -> bool:
     try:
         task = task_lib.Task.from_yaml(task_yaml)
         backend = cloud_vm_ray_backend.CloudVmRayBackend()
+        # Need to re-construct storage object in the controller process
+        # because when SkyPilot API server machine sends the yaml config to the
+        # controller machine, only storage metadata is sent, not the storage
+        # object itself.
+        for storage in task.storage_mounts.values():
+            storage.construct()
         backend.teardown_ephemeral_storage(task)
     except Exception as e:  # pylint: disable=broad-except
         logger.error('Failed to clean up storage: '
@@ -116,15 +122,17 @@ def _cleanup(service_name: str) -> bool:
             logger.error(f'Replica {info.replica_id} failed to terminate.')
     versions = serve_state.get_service_versions(service_name)
     serve_state.remove_service_versions(service_name)
-    success = True
-    for version in versions:
+
+    def cleanup_version_storage(version: int) -> bool:
         task_yaml: str = serve_utils.generate_task_yaml_file_name(
             service_name, version)
         logger.info(f'Cleaning up storage for version {version}, '
                     f'task_yaml: {task_yaml}')
-        success = success and cleanup_storage(task_yaml)
-    if not success:
+        return cleanup_storage(task_yaml)
+
+    if not all(map(cleanup_version_storage, versions)):
         failed = True
+
     return failed
 
 
@@ -148,7 +156,9 @@ def _start(service_name: str, tmp_task_yaml: str, job_id: int):
         controller_job_id=job_id,
         policy=service_spec.autoscaling_policy_str(),
         requested_resources_str=backend_utils.get_task_resources_str(task),
-        status=serve_state.ServiceStatus.CONTROLLER_INIT)
+        load_balancing_policy=service_spec.load_balancing_policy,
+        status=serve_state.ServiceStatus.CONTROLLER_INIT,
+        tls_encrypted=service_spec.tls_credential is not None)
     # Directly throw an error here. See sky/serve/api.py::up
     # for more details.
     if not success:
@@ -156,6 +166,10 @@ def _start(service_name: str, tmp_task_yaml: str, job_id: int):
         with ux_utils.print_exception_no_traceback():
             raise ValueError(f'Service {service_name} already exists.')
 
+    # Add initial version information to the service state.
+    serve_state.add_or_update_version(service_name, constants.INITIAL_VERSION,
+                                      service_spec)
+
     # Create the service working directory.
     service_dir = os.path.expanduser(
         serve_utils.generate_remote_service_dir_name(service_name))
@@ -182,19 +196,39 @@ def _start(service_name: str, tmp_task_yaml: str, job_id: int):
                 os.path.expanduser(constants.PORT_SELECTION_FILE_LOCK_PATH)):
             controller_port = common_utils.find_free_port(
                 constants.CONTROLLER_PORT_START)
+
+            # We expose the controller to the public network when running
+            # inside a kubernetes cluster to allow external load balancers
+            # (example, for high availability load balancers) to communicate
+            # with the controller.
+            def _get_host():
+                if 'KUBERNETES_SERVICE_HOST' in os.environ:
+                    return '0.0.0.0'
+                # Not using localhost to avoid using ipv6 address and causing
+                # the following error:
+                # ERROR:    [Errno 99] error while attempting to bind on address
+                # ('::1', 20001, 0, 0): cannot assign requested address
+                return '127.0.0.1'
+
+            controller_host = _get_host()
+
             # Start the controller.
             controller_process = multiprocessing.Process(
                 target=controller.run_controller,
-                args=(service_name, service_spec, task_yaml, controller_port))
+                args=(service_name, service_spec, task_yaml, controller_host,
+                      controller_port))
             controller_process.start()
             serve_state.set_service_controller_port(service_name,
                                                     controller_port)
 
-            # TODO(tian): Support HTTPS.
-            controller_addr = f'http://localhost:{controller_port}'
+            controller_addr = f'http://{controller_host}:{controller_port}'
+
             load_balancer_port = common_utils.find_free_port(
                 constants.LOAD_BALANCER_PORT_START)
 
+            # Extract the load balancing policy from the service spec
+            policy_name = service_spec.load_balancing_policy
+
             # Start the load balancer.
             # TODO(tian): Probably we could enable multiple ports specified in
             # service spec and we could start multiple load balancers.
@@ -203,7 +237,8 @@ def _start(service_name: str, tmp_task_yaml: str, job_id: int):
                 target=ux_utils.RedirectOutputForProcess(
                     load_balancer.run_load_balancer,
                     load_balancer_log_file).run,
-                args=(controller_addr, load_balancer_port))
+                args=(controller_addr, load_balancer_port, policy_name,
+                      service_spec.tls_credential))
             load_balancer_process.start()
             serve_state.set_service_load_balancer_port(service_name,
                                                        load_balancer_port)
@@ -215,15 +250,15 @@ def _start(service_name: str, tmp_task_yaml: str, job_id: int):
         serve_state.set_service_status_and_active_versions(
             service_name, serve_state.ServiceStatus.SHUTTING_DOWN)
     finally:
-        process_to_kill: List[multiprocessing.Process] = []
-        if load_balancer_process is not None:
-            process_to_kill.append(load_balancer_process)
-        if controller_process is not None:
-            process_to_kill.append(controller_process)
         # Kill load balancer process first since it will raise errors if failed
         # to connect to the controller. Then the controller process.
+        process_to_kill = [
+            proc for proc in [load_balancer_process, controller_process]
+            if proc is not None
+        ]
         subprocess_utils.kill_children_processes(
-            [process.pid for process in process_to_kill], force=True)
+            parent_pids=[process.pid for process in process_to_kill],
+            force=True)
         for process in process_to_kill:
             process.join()
         failed = _cleanup(service_name)
@@ -234,6 +269,7 @@ def _start(service_name: str, tmp_task_yaml: str, job_id: int):
         else:
             shutil.rmtree(service_dir)
             serve_state.remove_service(service_name)
+            serve_state.delete_all_versions(service_name)
             logger.info(f'Service {service_name} terminated successfully.')
 
 

sky/serve/service_spec.py

@@ -2,11 +2,14 @@
 import json
 import os
 import textwrap
-from typing import Any, Dict, Optional
+from typing import Any, Dict, List, Optional
 
 import yaml
 
+from sky import serve
 from sky.serve import constants
+from sky.serve import load_balancing_policies as lb_policies
+from sky.serve import serve_utils
 from sky.utils import common_utils
 from sky.utils import schemas
 from sky.utils import ux_utils
@@ -19,22 +22,19 @@ class SkyServiceSpec:
         self,
         readiness_path: str,
         initial_delay_seconds: int,
+        readiness_timeout_seconds: int,
         min_replicas: int,
         max_replicas: Optional[int] = None,
+        ports: Optional[str] = None,
         target_qps_per_replica: Optional[float] = None,
         post_data: Optional[Dict[str, Any]] = None,
+        tls_credential: Optional[serve_utils.TLSCredential] = None,
         readiness_headers: Optional[Dict[str, str]] = None,
         dynamic_ondemand_fallback: Optional[bool] = None,
         base_ondemand_fallback_replicas: Optional[int] = None,
         upscale_delay_seconds: Optional[int] = None,
         downscale_delay_seconds: Optional[int] = None,
-        # The following arguments are deprecated.
-        # TODO(ziming): remove this after 2 minor release, i.e. 0.6.0.
-        # Deprecated: Always be True
-        auto_restart: Optional[bool] = None,
-        # Deprecated: replaced by the target_qps_per_replica.
-        qps_upper_threshold: Optional[float] = None,
-        qps_lower_threshold: Optional[float] = None,
+        load_balancing_policy: Optional[str] = None,
     ) -> None:
         if max_replicas is not None and max_replicas < min_replicas:
             with ux_utils.print_exception_no_traceback():
@@ -61,27 +61,23 @@ class SkyServiceSpec:
                 raise ValueError('readiness_path must start with a slash (/). '
                                  f'Got: {readiness_path}')
 
-        # TODO(tian): Following field are deprecated. Remove after 2 minor
-        # release, i.e. 0.6.0.
-        if qps_upper_threshold is not None or qps_lower_threshold is not None:
+        # Add the check for unknown load balancing policies
+        if (load_balancing_policy is not None and
+                load_balancing_policy not in serve.LB_POLICIES):
             with ux_utils.print_exception_no_traceback():
                 raise ValueError(
-                    'Field `qps_upper_threshold` and `qps_lower_threshold`'
-                    'under `replica_policy` are deprecated. '
-                    'Please use target_qps_per_replica instead.')
-        if auto_restart is not None:
-            with ux_utils.print_exception_no_traceback():
-                raise ValueError(
-                    'Field `auto_restart` under `replica_policy` is deprecated.'
-                    'Currently, SkyServe will cleanup failed replicas'
-                    'and auto restart it to keep the service running.')
-
+                    f'Unknown load balancing policy: {load_balancing_policy}. '
+                    f'Available policies: {list(serve.LB_POLICIES.keys())}')
         self._readiness_path: str = readiness_path
         self._initial_delay_seconds: int = initial_delay_seconds
+        self._readiness_timeout_seconds: int = readiness_timeout_seconds
         self._min_replicas: int = min_replicas
         self._max_replicas: Optional[int] = max_replicas
+        self._ports: Optional[str] = ports
         self._target_qps_per_replica: Optional[float] = target_qps_per_replica
         self._post_data: Optional[Dict[str, Any]] = post_data
+        self._tls_credential: Optional[serve_utils.TLSCredential] = (
+            tls_credential)
         self._readiness_headers: Optional[Dict[str, str]] = readiness_headers
         self._dynamic_ondemand_fallback: Optional[
             bool] = dynamic_ondemand_fallback
@@ -89,6 +85,7 @@ class SkyServiceSpec:
             int] = base_ondemand_fallback_replicas
         self._upscale_delay_seconds: Optional[int] = upscale_delay_seconds
         self._downscale_delay_seconds: Optional[int] = downscale_delay_seconds
+        self._load_balancing_policy: Optional[str] = load_balancing_policy
 
         self._use_ondemand_fallback: bool = (
             self.dynamic_ondemand_fallback is not None and
@@ -113,16 +110,23 @@ class SkyServiceSpec:
             service_config['readiness_path'] = readiness_section
             initial_delay_seconds = None
             post_data = None
+            readiness_timeout_seconds = None
             readiness_headers = None
         else:
             service_config['readiness_path'] = readiness_section['path']
             initial_delay_seconds = readiness_section.get(
                 'initial_delay_seconds', None)
             post_data = readiness_section.get('post_data', None)
+            readiness_timeout_seconds = readiness_section.get(
+                'timeout_seconds', None)
             readiness_headers = readiness_section.get('headers', None)
         if initial_delay_seconds is None:
             initial_delay_seconds = constants.DEFAULT_INITIAL_DELAY_SECONDS
         service_config['initial_delay_seconds'] = initial_delay_seconds
+        if readiness_timeout_seconds is None:
+            readiness_timeout_seconds = (
+                constants.DEFAULT_READINESS_PROBE_TIMEOUT_SECONDS)
+        service_config['readiness_timeout_seconds'] = readiness_timeout_seconds
         if isinstance(post_data, str):
             try:
                 post_data = json.loads(post_data)
@@ -135,6 +139,14 @@ class SkyServiceSpec:
         service_config['post_data'] = post_data
         service_config['readiness_headers'] = readiness_headers
 
+        ports = config.get('ports', None)
+        if ports is not None:
+            assert isinstance(ports, int)
+            if not 1 <= ports <= 65535:
+                with ux_utils.print_exception_no_traceback():
+                    raise ValueError('Port must be between 1 and 65535.')
+        service_config['ports'] = str(ports) if ports is not None else None
+
         policy_section = config.get('replica_policy', None)
         simplified_policy_section = config.get('replicas', None)
         if policy_section is None or simplified_policy_section is not None:
@@ -151,14 +163,8 @@ class SkyServiceSpec:
             service_config['min_replicas'] = policy_section['min_replicas']
             service_config['max_replicas'] = policy_section.get(
                 'max_replicas', None)
-            service_config['qps_upper_threshold'] = policy_section.get(
-                'qps_upper_threshold', None)
-            service_config['qps_lower_threshold'] = policy_section.get(
-                'qps_lower_threshold', None)
             service_config['target_qps_per_replica'] = policy_section.get(
                 'target_qps_per_replica', None)
-            service_config['auto_restart'] = policy_section.get(
-                'auto_restart', None)
             service_config['upscale_delay_seconds'] = policy_section.get(
                 'upscale_delay_seconds', None)
             service_config['downscale_delay_seconds'] = policy_section.get(
@@ -169,6 +175,16 @@ class SkyServiceSpec:
             service_config['dynamic_ondemand_fallback'] = policy_section.get(
                 'dynamic_ondemand_fallback', None)
 
+        service_config['load_balancing_policy'] = config.get(
+            'load_balancing_policy', None)
+
+        tls_section = config.get('tls', None)
+        if tls_section is not None:
+            service_config['tls_credential'] = serve_utils.TLSCredential(
+                keyfile=tls_section.get('keyfile', None),
+                certfile=tls_section.get('certfile', None),
+            )
+
         return SkyServiceSpec(**service_config)
 
     @staticmethod
@@ -192,9 +208,12 @@ class SkyServiceSpec:
         return SkyServiceSpec.from_yaml_config(config['service'])
 
     def to_yaml_config(self) -> Dict[str, Any]:
-        config = dict()
+        config: Dict[str, Any] = {}
 
-        def add_if_not_none(section, key, value, no_empty: bool = False):
+        def add_if_not_none(section: str,
+                            key: Optional[str],
+                            value: Any,
+                            no_empty: bool = False):
             if no_empty and not value:
                 return
             if value is not None:
@@ -209,6 +228,8 @@ class SkyServiceSpec:
         add_if_not_none('readiness_probe', 'initial_delay_seconds',
                         self.initial_delay_seconds)
         add_if_not_none('readiness_probe', 'post_data', self.post_data)
+        add_if_not_none('readiness_probe', 'timeout_seconds',
+                        self.readiness_timeout_seconds)
         add_if_not_none('readiness_probe', 'headers', self._readiness_headers)
         add_if_not_none('replica_policy', 'min_replicas', self.min_replicas)
         add_if_not_none('replica_policy', 'max_replicas', self.max_replicas)
@@ -222,6 +243,12 @@ class SkyServiceSpec:
                         self.upscale_delay_seconds)
         add_if_not_none('replica_policy', 'downscale_delay_seconds',
                         self.downscale_delay_seconds)
+        add_if_not_none('load_balancing_policy', None,
+                        self._load_balancing_policy)
+        add_if_not_none('ports', None, int(self.ports) if self.ports else None)
+        if self.tls_credential is not None:
+            add_if_not_none('tls', 'keyfile', self.tls_credential.keyfile)
+            add_if_not_none('tls', 'certfile', self.tls_credential.certfile)
         return config
 
     def probe_str(self):
@@ -233,8 +260,8 @@ class SkyServiceSpec:
                    ' with custom headers')
         return f'{method}{headers}'
 
-    def spot_policy_str(self):
-        policy_strs = []
+    def spot_policy_str(self) -> str:
+        policy_strs: List[str] = []
         if (self.dynamic_ondemand_fallback is not None and
                 self.dynamic_ondemand_fallback):
             policy_strs.append('Dynamic on-demand fallback')
@@ -249,7 +276,9 @@ class SkyServiceSpec:
                 policy_strs.append('Static spot mixture with '
                                    f'{self.base_ondemand_fallback_replicas} '
                                    f'base on-demand replica{plural}')
-        return ' '.join(policy_strs) if policy_strs else 'No spot policy'
+        if not policy_strs:
+            return 'No spot fallback policy'
+        return ' '.join(policy_strs)
 
     def autoscaling_policy_str(self):
         # TODO(MaoZiming): Update policy_str
@@ -264,12 +293,24 @@ class SkyServiceSpec:
                 f'replica{max_plural} (target QPS per replica: '
                 f'{self.target_qps_per_replica})')
 
+    def set_ports(self, ports: str) -> None:
+        self._ports = ports
+
+    def tls_str(self):
+        if self.tls_credential is None:
+            return 'No TLS Enabled'
+        return (f'Keyfile: {self.tls_credential.keyfile}, '
+                f'Certfile: {self.tls_credential.certfile}')
+
     def __repr__(self) -> str:
         return textwrap.dedent(f"""\
             Readiness probe method:           {self.probe_str()}
             Readiness initial delay seconds:  {self.initial_delay_seconds}
+            Readiness probe timeout seconds:  {self.readiness_timeout_seconds}
             Replica autoscaling policy:       {self.autoscaling_policy_str()}
+            TLS Certificates:                 {self.tls_str()}
             Spot Policy:                      {self.spot_policy_str()}
+            Load Balancing Policy:            {self.load_balancing_policy}
         """)
 
     @property
@@ -280,6 +321,10 @@ class SkyServiceSpec:
     def initial_delay_seconds(self) -> int:
         return self._initial_delay_seconds
 
+    @property
+    def readiness_timeout_seconds(self) -> int:
+        return self._readiness_timeout_seconds
+
     @property
     def min_replicas(self) -> int:
         return self._min_replicas
@@ -289,6 +334,10 @@ class SkyServiceSpec:
         # If None, treated as having the same value of min_replicas.
         return self._max_replicas
 
+    @property
+    def ports(self) -> Optional[str]:
+        return self._ports
+
     @property
     def target_qps_per_replica(self) -> Optional[float]:
         return self._target_qps_per_replica
@@ -297,6 +346,15 @@ class SkyServiceSpec:
     def post_data(self) -> Optional[Dict[str, Any]]:
         return self._post_data
 
+    @property
+    def tls_credential(self) -> Optional[serve_utils.TLSCredential]:
+        return self._tls_credential
+
+    @tls_credential.setter
+    def tls_credential(self,
+                       value: Optional[serve_utils.TLSCredential]) -> None:
+        self._tls_credential = value
+
     @property
     def readiness_headers(self) -> Optional[Dict[str, str]]:
         return self._readiness_headers
@@ -320,3 +378,8 @@ class SkyServiceSpec:
     @property
     def use_ondemand_fallback(self) -> bool:
         return self._use_ondemand_fallback
+
+    @property
+    def load_balancing_policy(self) -> str:
+        return lb_policies.LoadBalancingPolicy.make_policy_name(
+            self._load_balancing_policy)

sky/server/__init__.py

@@ -0,0 +1 @@
+"""SkyPilot API Server."""