skypilot-nightly 1.0.0.dev20241220__py3-none-any.whl → 1.0.0.dev20241221__py3-none-any.whl

sky/__init__.py

@@ -5,7 +5,7 @@ from typing import Optional
 import urllib.request
 
 # Replaced with the current commit when building the wheels.
-_SKYPILOT_COMMIT_SHA = 'a73a9cb37152c5f05d06bd1afe29a685ef86f85c'
+_SKYPILOT_COMMIT_SHA = 'ee3cabd57247ff0f25cb65c0ee46bd35ead8d11a'
 
 
 def _get_git_commit():
@@ -35,7 +35,7 @@ def _get_git_commit():
 
 
 __commit__ = _get_git_commit()
-__version__ = '1.0.0.dev20241220'
+__version__ = '1.0.0.dev20241221'
 __root_dir__ = os.path.dirname(os.path.abspath(__file__))
 
 

sky/provision/kubernetes/utils.py

@@ -1234,7 +1234,8 @@ def construct_ssh_jump_command(
                               '-o StrictHostKeyChecking=no '
                               '-o UserKnownHostsFile=/dev/null '
                               f'-o IdentitiesOnly=yes '
-                              f'-W %h:%p {ssh_jump_user}@{ssh_jump_ip}')
+                              r'-W \[%h\]:%p '
+                              f'{ssh_jump_user}@{ssh_jump_ip}')
     if ssh_jump_port is not None:
         ssh_jump_proxy_command += f' -p {ssh_jump_port} '
     if proxy_cmd_path is not None:

sky/serve/core.py

@@ -384,6 +384,17 @@ def update(
         with ux_utils.print_exception_no_traceback():
             raise RuntimeError(prompt)
 
+    original_lb_policy = service_record['load_balancing_policy']
+    assert task.service is not None, 'Service section not found.'
+    if original_lb_policy != task.service.load_balancing_policy:
+        logger.warning(
+            f'{colorama.Fore.YELLOW}Current load balancing policy '
+            f'{original_lb_policy!r} is different from the new policy '
+            f'{task.service.load_balancing_policy!r}. Updating the load '
+            'balancing policy is not supported yet and it will be ignored. '
+            'The service will continue to use the current load balancing '
+            f'policy.{colorama.Style.RESET_ALL}')
+
     with rich_utils.safe_status(
             ux_utils.spinner_message('Initializing service')):
         controller_utils.maybe_translate_local_file_mounts_and_sync_up(
@@ -581,9 +592,10 @@ def status(
             'status': (sky.ServiceStatus) service status,
             'controller_port': (Optional[int]) controller port,
             'load_balancer_port': (Optional[int]) load balancer port,
-            'policy': (Optional[str]) load balancer policy description,
+            'policy': (Optional[str]) autoscaling policy description,
             'requested_resources_str': (str) str representation of
               requested resources,
+            'load_balancing_policy': (str) load balancing policy name,
             'replica_info': (List[Dict[str, Any]]) replica information,
         }
 

sky/serve/load_balancer.py

@@ -45,6 +45,8 @@ class SkyServeLoadBalancer:
         # Use the registry to create the load balancing policy
         self._load_balancing_policy = lb_policies.LoadBalancingPolicy.make(
             load_balancing_policy_name)
+        logger.info('Starting load balancer with policy '
+                    f'{load_balancing_policy_name}.')
         self._request_aggregator: serve_utils.RequestsAggregator = (
             serve_utils.RequestTimestamp())
         # TODO(tian): httpx.Client has a resource limit of 100 max connections
@@ -128,6 +130,7 @@ class SkyServeLoadBalancer:
             encountered if anything goes wrong.
         """
         logger.info(f'Proxy request to {url}')
+        self._load_balancing_policy.pre_execute_hook(url, request)
         try:
             # We defer the get of the client here on purpose, for case when the
             # replica is ready in `_proxy_with_retries` but refreshed before
@@ -147,11 +150,16 @@ class SkyServeLoadBalancer:
                 content=await request.body(),
                 timeout=constants.LB_STREAM_TIMEOUT)
             proxy_response = await client.send(proxy_request, stream=True)
+
+            async def background_func():
+                await proxy_response.aclose()
+                self._load_balancing_policy.post_execute_hook(url, request)
+
             return fastapi.responses.StreamingResponse(
                 content=proxy_response.aiter_raw(),
                 status_code=proxy_response.status_code,
                 headers=proxy_response.headers,
-                background=background.BackgroundTask(proxy_response.aclose))
+                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)}')
@@ -263,7 +271,7 @@ if __name__ == '__main__':
     parser.add_argument(
         '--load-balancing-policy',
         choices=available_policies,
-        default='round_robin',
+        default=lb_policies.DEFAULT_LB_POLICY,
         help=f'The load balancing policy to use. Available policies: '
         f'{", ".join(available_policies)}.')
     args = parser.parse_args()

sky/serve/load_balancing_policies.py

@@ -1,7 +1,9 @@
 """LoadBalancingPolicy: Policy to select endpoint."""
+import collections
 import random
+import threading
 import typing
-from typing import List, Optional
+from typing import Dict, List, Optional
 
 from sky import sky_logging
 
@@ -13,6 +15,10 @@ logger = sky_logging.init_logger(__name__)
 # Define a registry for load balancing policies
 LB_POLICIES = {}
 DEFAULT_LB_POLICY = None
+# Prior to #4439, the default policy was round_robin. We store the legacy
+# default policy here to maintain backwards compatibility. Remove this after
+# 2 minor release, i.e., 0.9.0.
+LEGACY_DEFAULT_POLICY = 'round_robin'
 
 
 def _request_repr(request: 'fastapi.Request') -> str:
@@ -38,11 +44,17 @@ class LoadBalancingPolicy:
             DEFAULT_LB_POLICY = name
 
     @classmethod
-    def make(cls, policy_name: Optional[str] = None) -> 'LoadBalancingPolicy':
-        """Create a load balancing policy from a name."""
+    def make_policy_name(cls, policy_name: Optional[str]) -> str:
+        """Return the policy name."""
+        assert DEFAULT_LB_POLICY is not None, 'No default policy set.'
         if policy_name is None:
-            policy_name = DEFAULT_LB_POLICY
+            return DEFAULT_LB_POLICY
+        return policy_name
 
+    @classmethod
+    def make(cls, policy_name: Optional[str] = None) -> 'LoadBalancingPolicy':
+        """Create a load balancing policy from a name."""
+        policy_name = cls.make_policy_name(policy_name)
         if policy_name not in LB_POLICIES:
             raise ValueError(f'Unknown load balancing policy: {policy_name}')
         return LB_POLICIES[policy_name]()
@@ -65,8 +77,16 @@ class LoadBalancingPolicy:
     def _select_replica(self, request: 'fastapi.Request') -> Optional[str]:
         raise NotImplementedError
 
+    def pre_execute_hook(self, replica_url: str,
+                         request: 'fastapi.Request') -> None:
+        pass
+
+    def post_execute_hook(self, replica_url: str,
+                          request: 'fastapi.Request') -> None:
+        pass
+
 
-class RoundRobinPolicy(LoadBalancingPolicy, name='round_robin', default=True):
+class RoundRobinPolicy(LoadBalancingPolicy, name='round_robin'):
     """Round-robin load balancing policy."""
 
     def __init__(self) -> None:
@@ -90,3 +110,43 @@ class RoundRobinPolicy(LoadBalancingPolicy, name='round_robin', default=True):
         ready_replica_url = self.ready_replicas[self.index]
         self.index = (self.index + 1) % len(self.ready_replicas)
         return ready_replica_url
+
+
+class LeastLoadPolicy(LoadBalancingPolicy, name='least_load', default=True):
+    """Least load load balancing policy."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.load_map: Dict[str, int] = collections.defaultdict(int)
+        self.lock = threading.Lock()
+
+    def set_ready_replicas(self, ready_replicas: List[str]) -> None:
+        if set(self.ready_replicas) == set(ready_replicas):
+            return
+        with self.lock:
+            self.ready_replicas = ready_replicas
+            for r in self.ready_replicas:
+                if r not in ready_replicas:
+                    del self.load_map[r]
+            for replica in ready_replicas:
+                self.load_map[replica] = self.load_map.get(replica, 0)
+
+    def _select_replica(self, request: 'fastapi.Request') -> Optional[str]:
+        del request  # Unused.
+        if not self.ready_replicas:
+            return None
+        with self.lock:
+            return min(self.ready_replicas,
+                       key=lambda replica: self.load_map.get(replica, 0))
+
+    def pre_execute_hook(self, replica_url: str,
+                         request: 'fastapi.Request') -> None:
+        del request  # Unused.
+        with self.lock:
+            self.load_map[replica_url] += 1
+
+    def post_execute_hook(self, replica_url: str,
+                          request: 'fastapi.Request') -> None:
+        del request  # Unused.
+        with self.lock:
+            self.load_map[replica_url] -= 1

sky/serve/serve_state.py

@@ -11,6 +11,7 @@ from typing import Any, Dict, List, Optional, Tuple
 import colorama
 
 from sky.serve import constants
+from sky.serve import load_balancing_policies as lb_policies
 from sky.utils import db_utils
 
 if typing.TYPE_CHECKING:
@@ -76,6 +77,8 @@ db_utils.add_column_to_table(_DB.cursor, _DB.conn, 'services',
 db_utils.add_column_to_table(_DB.cursor, _DB.conn, 'services',
                              'active_versions',
                              f'TEXT DEFAULT {json.dumps([])!r}')
+db_utils.add_column_to_table(_DB.cursor, _DB.conn, 'services',
+                             'load_balancing_policy', 'TEXT DEFAULT NULL')
 _UNIQUE_CONSTRAINT_FAILED_ERROR_MSG = 'UNIQUE constraint failed: services.name'
 
 
@@ -241,7 +244,8 @@ _SERVICE_STATUS_TO_COLOR = {
 
 
 def add_service(name: str, controller_job_id: int, policy: str,
-                requested_resources_str: str, status: ServiceStatus) -> bool:
+                requested_resources_str: str, load_balancing_policy: str,
+                status: ServiceStatus) -> bool:
     """Add a service in the database.
 
     Returns:
@@ -254,10 +258,10 @@ def add_service(name: str, controller_job_id: int, policy: str,
                 """\
                 INSERT INTO services
                 (name, controller_job_id, status, policy,
-                requested_resources_str)
-                VALUES (?, ?, ?, ?, ?)""",
+                requested_resources_str, load_balancing_policy)
+                VALUES (?, ?, ?, ?, ?, ?)""",
                 (name, controller_job_id, status.value, policy,
-                 requested_resources_str))
+                 requested_resources_str, load_balancing_policy))
 
     except sqlite3.IntegrityError as e:
         if str(e) != _UNIQUE_CONSTRAINT_FAILED_ERROR_MSG:
@@ -324,7 +328,12 @@ 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) = row[:13]
+     _, active_versions, load_balancing_policy) = row[:14]
+    if load_balancing_policy is None:
+        # This entry in database was added in #4439, and it will always be set
+        # to a str value. If it is None, it means it is an legacy entry and is
+        # using the legacy default policy.
+        load_balancing_policy = lb_policies.LEGACY_DEFAULT_POLICY
     return {
         'name': name,
         'controller_job_id': controller_job_id,
@@ -341,6 +350,7 @@ def _get_service_from_row(row) -> Dict[str, Any]:
         # integers in json format. This is mainly for display purpose.
         'active_versions': json.loads(active_versions),
         'requested_resources_str': requested_resources_str,
+        'load_balancing_policy': load_balancing_policy,
     }
 
 

sky/serve/serve_utils.py

@@ -811,7 +811,9 @@ 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 = []
@@ -832,6 +834,7 @@ def format_service_table(service_records: List[Dict[str, Any]],
         endpoint = get_endpoint(record)
         policy = record['policy']
         requested_resources_str = record['requested_resources_str']
+        load_balancing_policy = record['load_balancing_policy']
 
         service_values = [
             service_name,
@@ -842,7 +845,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)

sky/serve/service.py

@@ -150,6 +150,7 @@ 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),
+        load_balancing_policy=service_spec.load_balancing_policy,
         status=serve_state.ServiceStatus.CONTROLLER_INIT)
     # Directly throw an error here. See sky/serve/api.py::up
     # for more details.

sky/serve/service_spec.py

@@ -8,6 +8,7 @@ import yaml
 
 from sky import serve
 from sky.serve import constants
+from sky.serve import load_balancing_policies as lb_policies
 from sky.utils import common_utils
 from sky.utils import schemas
 from sky.utils import ux_utils
@@ -327,5 +328,6 @@ class SkyServiceSpec:
         return self._use_ondemand_fallback
 
     @property
-    def load_balancing_policy(self) -> Optional[str]:
-        return self._load_balancing_policy
+    def load_balancing_policy(self) -> str:
+        return lb_policies.LoadBalancingPolicy.make_policy_name(
+            self._load_balancing_policy)

{skypilot_nightly-1.0.0.dev20241220.dist-info → skypilot_nightly-1.0.0.dev20241221.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: skypilot-nightly
-Version: 1.0.0.dev20241220
+Version: 1.0.0.dev20241221
 Summary: SkyPilot: An intercloud broker for the clouds
 Author: SkyPilot Team
 License: Apache 2.0

{skypilot_nightly-1.0.0.dev20241220.dist-info → skypilot_nightly-1.0.0.dev20241221.dist-info}/RECORD

@@ -1,4 +1,4 @@
-sky/__init__.py,sha256=v_p7xNUZF9frwvqSyWlYtzPH1cANA1dA83B7dzndVl0,5944
+sky/__init__.py,sha256=PT4n67F3bRiUrj8lUmXPGOp1PG_ib48eEcF1y7gYmKQ,5944
 sky/admin_policy.py,sha256=hPo02f_A32gCqhUueF0QYy1fMSSKqRwYEg_9FxScN_s,3248
 sky/authentication.py,sha256=kACHmiZgWgRpYd1wx1ofbXRMErfMcFmWrkw4a9NxYrY,20988
 sky/check.py,sha256=Gg_62eCyNNYxu6oFiQTMr4XXJf7MjfZMdYOVS-v0SOc,9465
@@ -140,7 +140,7 @@ sky/provision/kubernetes/config.py,sha256=WEKcFXXhe89bLGAvoMiBvTDxdxkpTIA6ezrj2v
 sky/provision/kubernetes/instance.py,sha256=2pju1RDzH5vOXbOTgPzCToGdCjmqIouhA0wRUFytLnc,50257
 sky/provision/kubernetes/network.py,sha256=EpNjRQ131CXepqbdkoRKFu4szVrm0oKEpv1l8EgOkjU,12364
 sky/provision/kubernetes/network_utils.py,sha256=t1FS3K400fetH7cBuRgQJZl5_jEeMshsvsYmnMUcq8k,11399
-sky/provision/kubernetes/utils.py,sha256=PZ1hrgmzZnwtGGLNrT6PXG8Ke9geNGpdEr1OUew8sqQ,102869
+sky/provision/kubernetes/utils.py,sha256=W5CMFkkU3QEB219yPUZNOulsNghj_gcPJC6DOlZBwGM,102907
 sky/provision/kubernetes/manifests/smarter-device-manager-configmap.yaml,sha256=AMzYzlY0JIlfBWj5eX054Rc1XDW2thUcLSOGMJVhIdA,229
 sky/provision/kubernetes/manifests/smarter-device-manager-daemonset.yaml,sha256=RtTq4F1QUmR2Uunb6zuuRaPhV7hpesz4saHjn3Ncsb4,2010
 sky/provision/lambda_cloud/__init__.py,sha256=6EEvSgtUeEiup9ivIFevHmgv0GqleroO2X0K7TRa2nE,612
@@ -182,14 +182,14 @@ sky/serve/__init__.py,sha256=Bqw8nB9u1QF3ryjbV797SPZq0DWAcjT94E_5B8J24ag,1808
 sky/serve/autoscalers.py,sha256=khY1oZ22PRaUQNsLCoNKH178X_NiJw0LSLOKr7_LNgY,30275
 sky/serve/constants.py,sha256=7MflfgTHO9gDSux93U4BmNeEMWXxZB4q7I54KUwgp-s,4651
 sky/serve/controller.py,sha256=R5iIEGEEFtbm_6MvSGelYZP-vSmW0cSFuy64OexUc4g,11719
-sky/serve/core.py,sha256=gSf5jsRKnmHvTBE28zYqQ6xOYzYMuO4prYz9boSLUBs,31602
-sky/serve/load_balancer.py,sha256=I4W66eh1t1kA_C_VaMPI76WeDTCl3Z6rFxF6rQIWd6E,12636
-sky/serve/load_balancing_policies.py,sha256=_k4tkwIvhulR02Ln9ixYB_b97KOypr2xfSjMx8_zky0,3143
+sky/serve/core.py,sha256=YFwbow_36QJ6TE4POJhkDipmilULMY-HVCXTWo0MWyM,32308
+sky/serve/load_balancer.py,sha256=nNvDPJPRIrBc_qsBYJz1zzKa_fXDgfi0VDUf4SJEuW8,12990
+sky/serve/load_balancing_policies.py,sha256=XVj76qBgqh7h6wfx53RKQFzBefDWTE4TCdCEtFLLtI4,5398
 sky/serve/replica_managers.py,sha256=1xYDK9Te5wFEF5hUK0gyNIUib0MY-HScLHUBDlTSl-k,57774
-sky/serve/serve_state.py,sha256=Q7De4GoBEPxlN_t1Lpn-Y1fd94SeHZ3E-94f1OTuhpc,19086
-sky/serve/serve_utils.py,sha256=FYJdXaNvoMJFFt7x5Jua9PT-hLRL9J9nKI1BkkBtdxw,39452
-sky/serve/service.py,sha256=gVem2vX8XuR_1wTqwrzbszQAbjzjDP2ddd787aynT9g,12017
-sky/serve/service_spec.py,sha256=34dMQ37INHltBzWaxHl3y_o9X3wLOCWA5jUhmhH1II4,14740
+sky/serve/serve_state.py,sha256=Q274Ttm1oQ9blMVclyqyPmBONsh7oIAsIow_Q9imzjk,19762
+sky/serve/serve_utils.py,sha256=jmwpx6Oe6RL_VzLrp90yQffBFfMImT0cplvMvd0UNOE,39615
+sky/serve/service.py,sha256=7bvK9R9D48PZSYcOKSievXQ2mHUMk1d3AAIxtra7WOI,12083
+sky/serve/service_spec.py,sha256=2NHs4h_Bf8YcQ2LiOCDV47vMCm3pUmGMlUzVAwk_Qc0,14854
 sky/setup_files/MANIFEST.in,sha256=WF0T89NLichHxZDDSQzvSpiONtAEFyur2MPmGczgTIo,555
 sky/setup_files/dependencies.py,sha256=-_9ekkrmrU9vF7I9cqz_xP55qaSZ-qz2KfGVgQDQTTY,5903
 sky/setup_files/setup.py,sha256=HMqAIxHrhtQUOlm6_Iz5E_bL4dUvsYgXc9YVQIFayPs,7417
@@ -279,9 +279,9 @@ sky/utils/kubernetes/k8s_gpu_labeler_job.yaml,sha256=k0TBoQ4zgf79-sVkixKSGYFHQ7Z
 sky/utils/kubernetes/k8s_gpu_labeler_setup.yaml,sha256=VLKT2KKimZu1GDg_4AIlIt488oMQvhRZWwsj9vBbPUg,3812
 sky/utils/kubernetes/rsync_helper.sh,sha256=h4YwrPFf9727CACnMJvF3EyK_0OeOYKKt4su_daKekw,1256
 sky/utils/kubernetes/ssh_jump_lifecycle_manager.py,sha256=RFLJ3k7MR5UN4SKHykQ0lV9SgXumoULpKYIAt1vh-HU,6560
-skypilot_nightly-1.0.0.dev20241220.dist-info/LICENSE,sha256=emRJAvE7ngL6x0RhQvlns5wJzGI3NEQ_WMjNmd9TZc4,12170
-skypilot_nightly-1.0.0.dev20241220.dist-info/METADATA,sha256=uoKOU3FUyqPQueV1FoF8ied1VhOpu6GWNyaYAjCP6xw,20149
-skypilot_nightly-1.0.0.dev20241220.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
-skypilot_nightly-1.0.0.dev20241220.dist-info/entry_points.txt,sha256=StA6HYpuHj-Y61L2Ze-hK2IcLWgLZcML5gJu8cs6nU4,36
-skypilot_nightly-1.0.0.dev20241220.dist-info/top_level.txt,sha256=qA8QuiNNb6Y1OF-pCUtPEr6sLEwy2xJX06Bd_CrtrHY,4
-skypilot_nightly-1.0.0.dev20241220.dist-info/RECORD,,
+skypilot_nightly-1.0.0.dev20241221.dist-info/LICENSE,sha256=emRJAvE7ngL6x0RhQvlns5wJzGI3NEQ_WMjNmd9TZc4,12170
+skypilot_nightly-1.0.0.dev20241221.dist-info/METADATA,sha256=4844wOLUspN7oJE0chOWux3KrMtI_weLbd8tX7rYIeo,20149
+skypilot_nightly-1.0.0.dev20241221.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+skypilot_nightly-1.0.0.dev20241221.dist-info/entry_points.txt,sha256=StA6HYpuHj-Y61L2Ze-hK2IcLWgLZcML5gJu8cs6nU4,36
+skypilot_nightly-1.0.0.dev20241221.dist-info/top_level.txt,sha256=qA8QuiNNb6Y1OF-pCUtPEr6sLEwy2xJX06Bd_CrtrHY,4
+skypilot_nightly-1.0.0.dev20241221.dist-info/RECORD,,