skypilot-nightly 1.0.0.dev20250916__py3-none-any.whl → 1.0.0.dev20250919__py3-none-any.whl

sky/serve/serve_utils.py

@@ -745,8 +745,8 @@ def _get_service_status(
     return record
 
 
-def get_service_status_encoded(service_names: Optional[List[str]],
-                               pool: bool) -> str:
+def get_service_status_pickled(service_names: Optional[List[str]],
+                               pool: bool) -> List[Dict[str, str]]:
     service_statuses: List[Dict[str, str]] = []
     if service_names is None:
         # Get all service names
@@ -759,14 +759,34 @@ def get_service_status_encoded(service_names: Optional[List[str]],
             k: base64.b64encode(pickle.dumps(v)).decode('utf-8')
             for k, v in service_status.items()
         })
-    service_statuses = sorted(service_statuses, key=lambda x: x['name'])
+    return sorted(service_statuses, key=lambda x: x['name'])
+
+
+# TODO (kyuds): remove when serve codegen is removed
+def get_service_status_encoded(service_names: Optional[List[str]],
+                               pool: bool) -> str:
     # We have to use payload_type here to avoid the issue of
     # message_utils.decode_payload() not being able to correctly decode the
     # message with <sky-payload> tags.
+    service_statuses = get_service_status_pickled(service_names, pool)
     return message_utils.encode_payload(service_statuses,
                                         payload_type='service_status')
 
 
+def unpickle_service_status(
+        payload: List[Dict[str, str]]) -> List[Dict[str, Any]]:
+    service_statuses: List[Dict[str, Any]] = []
+    for service_status in payload:
+        if not isinstance(service_status, dict):
+            raise ValueError(f'Invalid service status: {service_status}')
+        service_statuses.append({
+            k: pickle.loads(base64.b64decode(v))
+            for k, v in service_status.items()
+        })
+    return service_statuses
+
+
+# TODO (kyuds): remove when serve codegen is removed
 def load_service_status(payload: str) -> List[Dict[str, Any]]:
     try:
         service_statuses_encoded = message_utils.decode_payload(
@@ -778,22 +798,16 @@ def load_service_status(payload: str) -> List[Dict[str, Any]]:
             service_statuses_encoded = message_utils.decode_payload(payload)
         else:
             raise
-    service_statuses: List[Dict[str, Any]] = []
-    for service_status in service_statuses_encoded:
-        if not isinstance(service_status, dict):
-            raise ValueError(f'Invalid service status: {service_status}')
-        service_statuses.append({
-            k: pickle.loads(base64.b64decode(v))
-            for k, v in service_status.items()
-        })
-    return service_statuses
+    return unpickle_service_status(service_statuses_encoded)
 
 
+# TODO (kyuds): remove when serve codegen is removed
 def add_version_encoded(service_name: str) -> str:
     new_version = serve_state.add_version(service_name)
     return message_utils.encode_payload(new_version)
 
 
+# TODO (kyuds): remove when serve codegen is removed
 def load_version_string(payload: str) -> str:
     return message_utils.decode_payload(payload)
 
@@ -999,6 +1013,8 @@ def wait_service_registration(service_name: str, job_id: int,
     Returns:
         Encoded load balancer port assigned to the service.
     """
+    # TODO (kyuds): when codegen is fully deprecated, return the lb port
+    # as an int directly instead of encoding it.
     start_time = time.time()
     setup_completed = False
     noun = 'pool' if pool else 'service'
@@ -1546,6 +1562,7 @@ def _format_replica_table(replica_records: List[Dict[str, Any]], show_all: bool,
 
 
 # =========================== CodeGen for Sky Serve ===========================
+# TODO (kyuds): deprecate and remove serve codegen entirely.
 
 
 # TODO(tian): Use REST API instead of SSH in the future. This codegen pattern

sky/serve/server/core.py

@@ -5,7 +5,9 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 from sky import backends
 from sky import exceptions
 from sky import sky_logging
+from sky.adaptors import common as adaptors_common
 from sky.backends import backend_utils
+from sky.serve import serve_rpc_utils
 from sky.serve import serve_utils
 from sky.serve.server import impl
 from sky.usage import usage_lib
@@ -13,7 +15,11 @@ from sky.utils import controller_utils
 from sky.utils import subprocess_utils
 
 if typing.TYPE_CHECKING:
+    import grpc
+
     import sky
+else:
+    grpc = adaptors_common.LazyImport('grpc')
 
 logger = sky_logging.init_logger(__name__)
 
@@ -105,25 +111,37 @@ def terminate_replica(service_name: str, replica_id: int, purge: bool) -> None:
         'Please spin up a service first.',
     )
 
-    backend = backend_utils.get_backend_from_handle(handle)
-    assert isinstance(backend, backends.CloudVmRayBackend)
-
-    code = serve_utils.ServeCodeGen.terminate_replica(service_name, replica_id,
-                                                      purge)
-    returncode, stdout, stderr = backend.run_on_head(handle,
-                                                     code,
-                                                     require_outputs=True,
-                                                     stream_logs=False,
-                                                     separate_stderr=True)
-
-    try:
-        subprocess_utils.handle_returncode(returncode,
-                                           code,
-                                           'Failed to terminate the replica',
-                                           stderr,
-                                           stream_logs=True)
-    except exceptions.CommandError as e:
-        raise RuntimeError(e.error_msg) from e
+    assert isinstance(handle, backends.CloudVmRayResourceHandle)
+    use_legacy = not handle.is_grpc_enabled_with_flag
+
+    if handle.is_grpc_enabled_with_flag:
+        try:
+            stdout = serve_rpc_utils.RpcRunner.terminate_replica(
+                handle, service_name, replica_id, purge)
+        except exceptions.SkyletMethodNotImplementedError:
+            use_legacy = True
+
+    if use_legacy:
+        backend = backend_utils.get_backend_from_handle(handle)
+        assert isinstance(backend, backends.CloudVmRayBackend)
+
+        code = serve_utils.ServeCodeGen.terminate_replica(
+            service_name, replica_id, purge)
+        returncode, stdout, stderr = backend.run_on_head(handle,
+                                                         code,
+                                                         require_outputs=True,
+                                                         stream_logs=False,
+                                                         separate_stderr=True)
+
+        try:
+            subprocess_utils.handle_returncode(
+                returncode,
+                code,
+                'Failed to terminate the replica',
+                stderr,
+                stream_logs=True)
+        except exceptions.CommandError as e:
+            raise RuntimeError(e.error_msg) from e
 
     sky_logging.print(stdout)
 

sky/serve/server/impl.py

@@ -5,6 +5,7 @@ import shlex
 import signal
 import tempfile
 import threading
+import typing
 from typing import Any, Dict, List, Optional, Set, Tuple, Union
 import uuid
 
@@ -17,10 +18,12 @@ from sky import execution
 from sky import sky_logging
 from sky import skypilot_config
 from sky import task as task_lib
+from sky.adaptors import common as adaptors_common
 from sky.backends import backend_utils
 from sky.catalog import common as service_catalog_common
 from sky.data import storage as storage_lib
 from sky.serve import constants as serve_constants
+from sky.serve import serve_rpc_utils
 from sky.serve import serve_state
 from sky.serve import serve_utils
 from sky.skylet import constants
@@ -36,6 +39,11 @@ from sky.utils import subprocess_utils
 from sky.utils import ux_utils
 from sky.utils import yaml_utils
 
+if typing.TYPE_CHECKING:
+    import grpc
+else:
+    grpc = adaptors_common.LazyImport('grpc')
+
 logger = sky_logging.init_logger(__name__)
 
 
@@ -78,24 +86,35 @@ def _get_service_record(
     """Get the service record."""
     noun = 'pool' if pool else 'service'
 
-    code = serve_utils.ServeCodeGen.get_service_status([service_name],
-                                                       pool=pool)
-    returncode, serve_status_payload, stderr = backend.run_on_head(
-        handle,
-        code,
-        require_outputs=True,
-        stream_logs=False,
-        separate_stderr=True)
-    try:
-        subprocess_utils.handle_returncode(returncode,
-                                           code,
-                                           f'Failed to get {noun} status',
-                                           stderr,
-                                           stream_logs=True)
-    except exceptions.CommandError as e:
-        raise RuntimeError(e.error_msg) from e
+    assert isinstance(handle, backends.CloudVmRayResourceHandle)
+    use_legacy = not handle.is_grpc_enabled_with_flag
 
-    service_statuses = serve_utils.load_service_status(serve_status_payload)
+    if handle.is_grpc_enabled_with_flag:
+        try:
+            service_statuses = serve_rpc_utils.RpcRunner.get_service_status(
+                handle, [service_name], pool)
+        except exceptions.SkyletMethodNotImplementedError:
+            use_legacy = True
+
+    if use_legacy:
+        code = serve_utils.ServeCodeGen.get_service_status([service_name],
+                                                           pool=pool)
+        returncode, serve_status_payload, stderr = backend.run_on_head(
+            handle,
+            code,
+            require_outputs=True,
+            stream_logs=False,
+            separate_stderr=True)
+        try:
+            subprocess_utils.handle_returncode(returncode,
+                                               code,
+                                               f'Failed to get {noun} status',
+                                               stderr,
+                                               stream_logs=True)
+        except exceptions.CommandError as e:
+            raise RuntimeError(e.error_msg) from e
+
+        service_statuses = serve_utils.load_service_status(serve_status_payload)
 
     assert len(service_statuses) <= 1, service_statuses
     if not service_statuses:
@@ -287,30 +306,44 @@ def up(
         fore = colorama.Fore
 
         assert controller_job_id is not None and controller_handle is not None
+        assert isinstance(controller_handle, backends.CloudVmRayResourceHandle)
+        backend = backend_utils.get_backend_from_handle(controller_handle)
+        assert isinstance(backend, backends.CloudVmRayBackend)
         # TODO(tian): Cache endpoint locally to speedup. Endpoint won't
         # change after the first time, so there is no consistency issue.
-        with rich_utils.safe_status(
-                ux_utils.spinner_message(
-                    f'Waiting for the {noun} to register')):
-            # This function will check the controller job id in the database
-            # 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, pool)
-            backend = backend_utils.get_backend_from_handle(controller_handle)
-            assert isinstance(backend, backends.CloudVmRayBackend)
-            assert isinstance(controller_handle,
-                              backends.CloudVmRayResourceHandle)
-            returncode, lb_port_payload, _ = backend.run_on_head(
-                controller_handle,
-                code,
-                require_outputs=True,
-                stream_logs=False)
         try:
-            subprocess_utils.handle_returncode(
-                returncode, code, f'Failed to wait for {noun} initialization',
-                lb_port_payload)
-        except exceptions.CommandError:
+            with rich_utils.safe_status(
+                    ux_utils.spinner_message(
+                        f'Waiting for the {noun} to register')):
+                # This function will check the controller job id in the database
+                # and return the endpoint if the job id matches. Otherwise it
+                # will return None.
+                use_legacy = not controller_handle.is_grpc_enabled_with_flag
+
+                if controller_handle.is_grpc_enabled_with_flag:
+                    try:
+                        lb_port = serve_rpc_utils.RpcRunner.wait_service_registration(  # pylint: disable=line-too-long
+                            controller_handle, service_name, controller_job_id,
+                            pool)
+                    except exceptions.SkyletMethodNotImplementedError:
+                        use_legacy = True
+
+                if use_legacy:
+                    code = serve_utils.ServeCodeGen.wait_service_registration(
+                        service_name, controller_job_id, pool)
+                    returncode, lb_port_payload, _ = backend.run_on_head(
+                        controller_handle,
+                        code,
+                        require_outputs=True,
+                        stream_logs=False)
+                    subprocess_utils.handle_returncode(
+                        returncode, code,
+                        f'Failed to wait for {noun} initialization',
+                        lb_port_payload)
+                    lb_port = serve_utils.load_service_initialization_result(
+                        lb_port_payload)
+        except (exceptions.CommandError, grpc.FutureTimeoutError,
+                grpc.RpcError):
             if serve_utils.is_consolidation_mode(pool):
                 with ux_utils.print_exception_no_traceback():
                     raise RuntimeError(
@@ -344,8 +377,6 @@ def up(
                         'Failed to spin up the service. Please '
                         'check the logs above for more details.') from None
         else:
-            lb_port = serve_utils.load_service_initialization_result(
-                lb_port_payload)
             if not serve_utils.is_consolidation_mode(pool) and not pool:
                 socket_endpoint = backend_utils.get_endpoints(
                     controller_handle.cluster_name,
@@ -461,6 +492,7 @@ def update(
         f'use {ux_utils.BOLD}sky serve up{ux_utils.RESET_BOLD}',
     )
 
+    assert isinstance(handle, backends.CloudVmRayResourceHandle)
     backend = backend_utils.get_backend_from_handle(handle)
     assert isinstance(backend, backends.CloudVmRayBackend)
 
@@ -503,29 +535,39 @@ def update(
         controller_utils.maybe_translate_local_file_mounts_and_sync_up(
             task, task_type='serve')
 
-    code = serve_utils.ServeCodeGen.add_version(service_name)
-    returncode, version_string_payload, stderr = backend.run_on_head(
-        handle,
-        code,
-        require_outputs=True,
-        stream_logs=False,
-        separate_stderr=True)
-    try:
-        subprocess_utils.handle_returncode(returncode,
-                                           code,
-                                           'Failed to add version',
-                                           stderr,
-                                           stream_logs=True)
-    except exceptions.CommandError as e:
-        raise RuntimeError(e.error_msg) from e
+    use_legacy = not handle.is_grpc_enabled_with_flag
 
-    version_string = serve_utils.load_version_string(version_string_payload)
-    try:
-        current_version = int(version_string)
-    except ValueError as e:
-        with ux_utils.print_exception_no_traceback():
-            raise ValueError(f'Failed to parse version: {version_string}; '
-                             f'Returncode: {returncode}') from e
+    if handle.is_grpc_enabled_with_flag:
+        try:
+            current_version = serve_rpc_utils.RpcRunner.add_version(
+                handle, service_name)
+        except exceptions.SkyletMethodNotImplementedError:
+            use_legacy = True
+
+    if use_legacy:
+        code = serve_utils.ServeCodeGen.add_version(service_name)
+        returncode, version_string_payload, stderr = backend.run_on_head(
+            handle,
+            code,
+            require_outputs=True,
+            stream_logs=False,
+            separate_stderr=True)
+        try:
+            subprocess_utils.handle_returncode(returncode,
+                                               code,
+                                               'Failed to add version',
+                                               stderr,
+                                               stream_logs=True)
+        except exceptions.CommandError as e:
+            raise RuntimeError(e.error_msg) from e
+
+        version_string = serve_utils.load_version_string(version_string_payload)
+        try:
+            current_version = int(version_string)
+        except ValueError as e:
+            with ux_utils.print_exception_no_traceback():
+                raise ValueError(f'Failed to parse version: {version_string}; '
+                                 f'Returncode: {returncode}') from e
 
     with tempfile.NamedTemporaryFile(
             prefix=f'{service_name}-v{current_version}',
@@ -540,23 +582,33 @@ def update(
                                      {remote_task_yaml_path: service_file.name},
                                      storage_mounts=None)
 
-        code = serve_utils.ServeCodeGen.update_service(service_name,
-                                                       current_version,
-                                                       mode=mode.value,
-                                                       pool=pool)
-        returncode, _, stderr = backend.run_on_head(handle,
-                                                    code,
-                                                    require_outputs=True,
-                                                    stream_logs=False,
-                                                    separate_stderr=True)
-        try:
-            subprocess_utils.handle_returncode(returncode,
-                                               code,
-                                               f'Failed to update {noun}s',
-                                               stderr,
-                                               stream_logs=True)
-        except exceptions.CommandError as e:
-            raise RuntimeError(e.error_msg) from e
+        use_legacy = not handle.is_grpc_enabled_with_flag
+
+        if handle.is_grpc_enabled_with_flag:
+            try:
+                serve_rpc_utils.RpcRunner.update_service(
+                    handle, service_name, current_version, mode, pool)
+            except exceptions.SkyletMethodNotImplementedError:
+                use_legacy = True
+
+        if use_legacy:
+            code = serve_utils.ServeCodeGen.update_service(service_name,
+                                                           current_version,
+                                                           mode=mode.value,
+                                                           pool=pool)
+            returncode, _, stderr = backend.run_on_head(handle,
+                                                        code,
+                                                        require_outputs=True,
+                                                        stream_logs=False,
+                                                        separate_stderr=True)
+            try:
+                subprocess_utils.handle_returncode(returncode,
+                                                   code,
+                                                   f'Failed to update {noun}s',
+                                                   stderr,
+                                                   stream_logs=True)
+            except exceptions.CommandError as e:
+                raise RuntimeError(e.error_msg) from e
 
     cmd = 'sky jobs pool status' if pool else 'sky serve status'
     logger.info(
@@ -619,29 +671,44 @@ def down(
         raise ValueError(f'Can only specify one of {noun}_names or all. '
                          f'Provided {argument_str!r}.')
 
-    backend = backend_utils.get_backend_from_handle(handle)
-    assert isinstance(backend, backends.CloudVmRayBackend)
     service_names = None if all else service_names
-    code = serve_utils.ServeCodeGen.terminate_services(service_names, purge,
-                                                       pool)
 
     try:
-        returncode, stdout, _ = backend.run_on_head(handle,
-                                                    code,
-                                                    require_outputs=True,
-                                                    stream_logs=False)
+        assert isinstance(handle, backends.CloudVmRayResourceHandle)
+        use_legacy = not handle.is_grpc_enabled_with_flag
+
+        if handle.is_grpc_enabled_with_flag:
+            try:
+                stdout = serve_rpc_utils.RpcRunner.terminate_services(
+                    handle, service_names, purge, pool)
+            except exceptions.SkyletMethodNotImplementedError:
+                use_legacy = True
+
+        if use_legacy:
+            backend = backend_utils.get_backend_from_handle(handle)
+            assert isinstance(backend, backends.CloudVmRayBackend)
+            code = serve_utils.ServeCodeGen.terminate_services(
+                service_names, purge, pool)
+
+            returncode, stdout, _ = backend.run_on_head(handle,
+                                                        code,
+                                                        require_outputs=True,
+                                                        stream_logs=False)
+
+            subprocess_utils.handle_returncode(returncode, code,
+                                               f'Failed to terminate {noun}',
+                                               stdout)
     except exceptions.FetchClusterInfoError as e:
         raise RuntimeError(
             'Failed to fetch controller IP. Please refresh controller status '
-            f'by `sky status -r {controller_type.value.cluster_name}` '
-            'and try again.') from e
-
-    try:
-        subprocess_utils.handle_returncode(returncode, code,
-                                           f'Failed to terminate {noun}',
-                                           stdout)
+            f'by `sky status -r {controller_type.value.cluster_name}` and try '
+            'again.') from e
     except exceptions.CommandError as e:
         raise RuntimeError(e.error_msg) from e
+    except grpc.RpcError as e:
+        raise RuntimeError(f'{e.details()} ({e.code()})') from e
+    except grpc.FutureTimeoutError as e:
+        raise RuntimeError('gRPC timed out') from e
 
     logger.info(stdout)
 
@@ -669,27 +736,40 @@ def status(
         stopped_message=controller_type.value.default_hint_if_non_existent.
         replace('service', noun))
 
-    backend = backend_utils.get_backend_from_handle(handle)
-    assert isinstance(backend, backends.CloudVmRayBackend)
+    assert isinstance(handle, backends.CloudVmRayResourceHandle)
+    use_legacy = not handle.is_grpc_enabled_with_flag
 
-    code = serve_utils.ServeCodeGen.get_service_status(service_names, pool=pool)
-    returncode, serve_status_payload, stderr = backend.run_on_head(
-        handle,
-        code,
-        require_outputs=True,
-        stream_logs=False,
-        separate_stderr=True)
+    if handle.is_grpc_enabled_with_flag:
+        try:
+            service_records = serve_rpc_utils.RpcRunner.get_service_status(
+                handle, service_names, pool)
+        except exceptions.SkyletMethodNotImplementedError:
+            use_legacy = True
+
+    if use_legacy:
+        backend = backend_utils.get_backend_from_handle(handle)
+        assert isinstance(backend, backends.CloudVmRayBackend)
+
+        code = serve_utils.ServeCodeGen.get_service_status(service_names,
+                                                           pool=pool)
+        returncode, serve_status_payload, stderr = backend.run_on_head(
+            handle,
+            code,
+            require_outputs=True,
+            stream_logs=False,
+            separate_stderr=True)
 
-    try:
-        subprocess_utils.handle_returncode(returncode,
-                                           code,
-                                           f'Failed to fetch {noun}s',
-                                           stderr,
-                                           stream_logs=True)
-    except exceptions.CommandError as e:
-        raise RuntimeError(e.error_msg) from e
+        try:
+            subprocess_utils.handle_returncode(returncode,
+                                               code,
+                                               f'Failed to fetch {noun}s',
+                                               stderr,
+                                               stream_logs=True)
+        except exceptions.CommandError as e:
+            raise RuntimeError(e.error_msg) from e
+
+        service_records = serve_utils.load_service_status(serve_status_payload)
 
-    service_records = serve_utils.load_service_status(serve_status_payload)
     # Get the endpoint for each service
     for service_record in service_records:
         service_record['endpoint'] = None
@@ -792,25 +872,37 @@ def _get_all_replica_targets(
         handle: backends.CloudVmRayResourceHandle,
         pool: bool) -> Set[serve_utils.ServiceComponentTarget]:
     """Helper function to get targets for all live replicas."""
-    code = serve_utils.ServeCodeGen.get_service_status([service_name],
-                                                       pool=pool)
-    returncode, serve_status_payload, stderr = backend.run_on_head(
-        handle,
-        code,
-        require_outputs=True,
-        stream_logs=False,
-        separate_stderr=True)
+    assert isinstance(handle, backends.CloudVmRayResourceHandle)
+    use_legacy = not handle.is_grpc_enabled_with_flag
 
-    try:
-        subprocess_utils.handle_returncode(returncode,
-                                           code,
-                                           'Failed to fetch services',
-                                           stderr,
-                                           stream_logs=True)
-    except exceptions.CommandError as e:
-        raise RuntimeError(e.error_msg) from e
+    if handle.is_grpc_enabled_with_flag:
+        try:
+            service_records = serve_rpc_utils.RpcRunner.get_service_status(
+                handle, [service_name], pool)
+        except exceptions.SkyletMethodNotImplementedError:
+            use_legacy = True
+
+    if use_legacy:
+        code = serve_utils.ServeCodeGen.get_service_status([service_name],
+                                                           pool=pool)
+        returncode, serve_status_payload, stderr = backend.run_on_head(
+            handle,
+            code,
+            require_outputs=True,
+            stream_logs=False,
+            separate_stderr=True)
+
+        try:
+            subprocess_utils.handle_returncode(returncode,
+                                               code,
+                                               'Failed to fetch services',
+                                               stderr,
+                                               stream_logs=True)
+        except exceptions.CommandError as e:
+            raise RuntimeError(e.error_msg) from e
+
+        service_records = serve_utils.load_service_status(serve_status_payload)
 
-    service_records = serve_utils.load_service_status(serve_status_payload)
     if not service_records:
         raise ValueError(f'Service {service_name!r} not found.')
     assert len(service_records) == 1