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

sky/schemas/generated/servev1_pb2_grpc.py

@@ -0,0 +1,322 @@
+# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+"""Client and server classes corresponding to protobuf-defined services."""
+import grpc
+import warnings
+
+from sky.schemas.generated import servev1_pb2 as sky_dot_schemas_dot_generated_dot_servev1__pb2
+
+GRPC_GENERATED_VERSION = '1.63.0'
+GRPC_VERSION = grpc.__version__
+EXPECTED_ERROR_RELEASE = '1.65.0'
+SCHEDULED_RELEASE_DATE = 'June 25, 2024'
+_version_not_supported = False
+
+try:
+    from grpc._utilities import first_version_is_lower
+    _version_not_supported = first_version_is_lower(GRPC_VERSION, GRPC_GENERATED_VERSION)
+except ImportError:
+    _version_not_supported = True
+
+if _version_not_supported:
+    warnings.warn(
+        f'The grpc package installed is at version {GRPC_VERSION},'
+        + f' but the generated code in sky/schemas/generated/servev1_pb2_grpc.py depends on'
+        + f' grpcio>={GRPC_GENERATED_VERSION}.'
+        + f' Please upgrade your grpc module to grpcio>={GRPC_GENERATED_VERSION}'
+        + f' or downgrade your generated code using grpcio-tools<={GRPC_VERSION}.'
+        + f' This warning will become an error in {EXPECTED_ERROR_RELEASE},'
+        + f' scheduled for release on {SCHEDULED_RELEASE_DATE}.',
+        RuntimeWarning
+    )
+
+
+class ServeServiceStub(object):
+    """Missing associated documentation comment in .proto file."""
+
+    def __init__(self, channel):
+        """Constructor.
+
+        Args:
+            channel: A grpc.Channel.
+        """
+        self.GetServiceStatus = channel.unary_unary(
+                '/serve.v1.ServeService/GetServiceStatus',
+                request_serializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.GetServiceStatusRequest.SerializeToString,
+                response_deserializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.GetServiceStatusResponse.FromString,
+                _registered_method=True)
+        self.AddVersion = channel.unary_unary(
+                '/serve.v1.ServeService/AddVersion',
+                request_serializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.AddVersionRequest.SerializeToString,
+                response_deserializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.AddVersionResponse.FromString,
+                _registered_method=True)
+        self.TerminateServices = channel.unary_unary(
+                '/serve.v1.ServeService/TerminateServices',
+                request_serializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.TerminateServicesRequest.SerializeToString,
+                response_deserializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.TerminateServicesResponse.FromString,
+                _registered_method=True)
+        self.TerminateReplica = channel.unary_unary(
+                '/serve.v1.ServeService/TerminateReplica',
+                request_serializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.TerminateReplicaRequest.SerializeToString,
+                response_deserializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.TerminateReplicaResponse.FromString,
+                _registered_method=True)
+        self.WaitServiceRegistration = channel.unary_unary(
+                '/serve.v1.ServeService/WaitServiceRegistration',
+                request_serializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.WaitServiceRegistrationRequest.SerializeToString,
+                response_deserializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.WaitServiceRegistrationResponse.FromString,
+                _registered_method=True)
+        self.UpdateService = channel.unary_unary(
+                '/serve.v1.ServeService/UpdateService',
+                request_serializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.UpdateServiceRequest.SerializeToString,
+                response_deserializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.UpdateServiceResponse.FromString,
+                _registered_method=True)
+
+
+class ServeServiceServicer(object):
+    """Missing associated documentation comment in .proto file."""
+
+    def GetServiceStatus(self, request, context):
+        """Get status of service.
+        """
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def AddVersion(self, request, context):
+        """Add version to service.
+        """
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def TerminateServices(self, request, context):
+        """Terminate services.
+        """
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def TerminateReplica(self, request, context):
+        """Terminate replica.
+        """
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def WaitServiceRegistration(self, request, context):
+        """Wait for service registration.
+        """
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+    def UpdateService(self, request, context):
+        """Update service.
+        """
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+
+def add_ServeServiceServicer_to_server(servicer, server):
+    rpc_method_handlers = {
+            'GetServiceStatus': grpc.unary_unary_rpc_method_handler(
+                    servicer.GetServiceStatus,
+                    request_deserializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.GetServiceStatusRequest.FromString,
+                    response_serializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.GetServiceStatusResponse.SerializeToString,
+            ),
+            'AddVersion': grpc.unary_unary_rpc_method_handler(
+                    servicer.AddVersion,
+                    request_deserializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.AddVersionRequest.FromString,
+                    response_serializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.AddVersionResponse.SerializeToString,
+            ),
+            'TerminateServices': grpc.unary_unary_rpc_method_handler(
+                    servicer.TerminateServices,
+                    request_deserializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.TerminateServicesRequest.FromString,
+                    response_serializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.TerminateServicesResponse.SerializeToString,
+            ),
+            'TerminateReplica': grpc.unary_unary_rpc_method_handler(
+                    servicer.TerminateReplica,
+                    request_deserializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.TerminateReplicaRequest.FromString,
+                    response_serializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.TerminateReplicaResponse.SerializeToString,
+            ),
+            'WaitServiceRegistration': grpc.unary_unary_rpc_method_handler(
+                    servicer.WaitServiceRegistration,
+                    request_deserializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.WaitServiceRegistrationRequest.FromString,
+                    response_serializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.WaitServiceRegistrationResponse.SerializeToString,
+            ),
+            'UpdateService': grpc.unary_unary_rpc_method_handler(
+                    servicer.UpdateService,
+                    request_deserializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.UpdateServiceRequest.FromString,
+                    response_serializer=sky_dot_schemas_dot_generated_dot_servev1__pb2.UpdateServiceResponse.SerializeToString,
+            ),
+    }
+    generic_handler = grpc.method_handlers_generic_handler(
+            'serve.v1.ServeService', rpc_method_handlers)
+    server.add_generic_rpc_handlers((generic_handler,))
+
+
+ # This class is part of an EXPERIMENTAL API.
+class ServeService(object):
+    """Missing associated documentation comment in .proto file."""
+
+    @staticmethod
+    def GetServiceStatus(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            '/serve.v1.ServeService/GetServiceStatus',
+            sky_dot_schemas_dot_generated_dot_servev1__pb2.GetServiceStatusRequest.SerializeToString,
+            sky_dot_schemas_dot_generated_dot_servev1__pb2.GetServiceStatusResponse.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)
+
+    @staticmethod
+    def AddVersion(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            '/serve.v1.ServeService/AddVersion',
+            sky_dot_schemas_dot_generated_dot_servev1__pb2.AddVersionRequest.SerializeToString,
+            sky_dot_schemas_dot_generated_dot_servev1__pb2.AddVersionResponse.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)
+
+    @staticmethod
+    def TerminateServices(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            '/serve.v1.ServeService/TerminateServices',
+            sky_dot_schemas_dot_generated_dot_servev1__pb2.TerminateServicesRequest.SerializeToString,
+            sky_dot_schemas_dot_generated_dot_servev1__pb2.TerminateServicesResponse.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)
+
+    @staticmethod
+    def TerminateReplica(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            '/serve.v1.ServeService/TerminateReplica',
+            sky_dot_schemas_dot_generated_dot_servev1__pb2.TerminateReplicaRequest.SerializeToString,
+            sky_dot_schemas_dot_generated_dot_servev1__pb2.TerminateReplicaResponse.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)
+
+    @staticmethod
+    def WaitServiceRegistration(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            '/serve.v1.ServeService/WaitServiceRegistration',
+            sky_dot_schemas_dot_generated_dot_servev1__pb2.WaitServiceRegistrationRequest.SerializeToString,
+            sky_dot_schemas_dot_generated_dot_servev1__pb2.WaitServiceRegistrationResponse.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)
+
+    @staticmethod
+    def UpdateService(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            '/serve.v1.ServeService/UpdateService',
+            sky_dot_schemas_dot_generated_dot_servev1__pb2.UpdateServiceRequest.SerializeToString,
+            sky_dot_schemas_dot_generated_dot_servev1__pb2.UpdateServiceResponse.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)

sky/serve/serve_rpc_utils.py

@@ -0,0 +1,179 @@
+"""Rpc Utilities for SkyServe"""
+
+import typing
+from typing import Any, Dict, List, Optional, Tuple
+
+from sky import backends
+from sky.adaptors import common as adaptors_common
+from sky.backends import backend_utils
+from sky.serve import serve_utils
+
+if typing.TYPE_CHECKING:
+    from sky.schemas.generated import servev1_pb2
+else:
+    servev1_pb2 = adaptors_common.LazyImport(
+        'sky.schemas.generated.servev1_pb2')
+
+# ======================= gRPC Converters for Sky Serve =======================
+
+
+class GetServiceStatusRequestConverter:
+    """Converter for GetServiceStatusRequest"""
+
+    @classmethod
+    def to_proto(cls, service_names: Optional[List[str]],
+                 pool: bool) -> 'servev1_pb2.GetServiceStatusRequest':
+        request = servev1_pb2.GetServiceStatusRequest()
+        request.pool = pool
+        if service_names is not None:
+            request.service_names.names.extend(service_names)
+        return request
+
+    @classmethod
+    def from_proto(
+        cls, proto: 'servev1_pb2.GetServiceStatusRequest'
+    ) -> Tuple[Optional[List[str]], bool]:
+        pool = proto.pool
+        if proto.HasField('service_names'):
+            service_names = list(proto.service_names.names)
+        else:
+            service_names = None
+        return service_names, pool
+
+
+class GetServiceStatusResponseConverter:
+    """Converter for GetServiceStatusResponse"""
+
+    @classmethod
+    def to_proto(
+        cls,
+        statuses: List[Dict[str,
+                            str]]) -> 'servev1_pb2.GetServiceStatusResponse':
+        response = servev1_pb2.GetServiceStatusResponse()
+        for status in statuses:
+            added = response.statuses.add()
+            added.status.update(status)
+        return response
+
+    @classmethod
+    def from_proto(
+            cls, proto: 'servev1_pb2.GetServiceStatusResponse'
+    ) -> List[Dict[str, str]]:
+        pickled = [dict(status.status) for status in proto.statuses]
+        return pickled
+
+
+class TerminateServicesRequestConverter:
+    """Converter for TerminateServicesRequest"""
+
+    @classmethod
+    def to_proto(cls, service_names: Optional[List[str]], purge: bool,
+                 pool: bool) -> 'servev1_pb2.TerminateServicesRequest':
+        request = servev1_pb2.TerminateServicesRequest()
+        request.purge = purge
+        request.pool = pool
+        if service_names is not None:
+            request.service_names.names.extend(service_names)
+        return request
+
+    @classmethod
+    def from_proto(
+        cls, proto: 'servev1_pb2.TerminateServicesRequest'
+    ) -> Tuple[Optional[List[str]], bool, bool]:
+        purge = proto.purge
+        pool = proto.pool
+        if proto.HasField('service_names'):
+            service_names = list(proto.service_names.names)
+        else:
+            service_names = None
+        return service_names, purge, pool
+
+
+# ========================= gRPC Runner for Sky Serve =========================
+
+
+class RpcRunner:
+    """gRPC Runner for Sky Serve
+
+    The RPC runner does not catch errors, and assumes that backend handle has
+    grpc enabled.
+
+    Common exceptions raised:
+        exceptions.FetchClusterInfoError
+        exceptions.SkyletInternalError
+        grpc.RpcError
+        grpc.FutureTimeoutError
+        AssertionError
+    """
+
+    @classmethod
+    def get_service_status(cls, handle: backends.CloudVmRayResourceHandle,
+                           service_names: Optional[List[str]],
+                           pool: bool) -> List[Dict[str, Any]]:
+        assert handle.is_grpc_enabled_with_flag
+        request = GetServiceStatusRequestConverter.to_proto(service_names, pool)
+        response = backend_utils.invoke_skylet_with_retries(
+            lambda: backends.SkyletClient(handle.get_grpc_channel()
+                                         ).get_service_status(request))
+        pickled = GetServiceStatusResponseConverter.from_proto(response)
+        return serve_utils.unpickle_service_status(pickled)
+
+    @classmethod
+    def add_version(cls, handle: backends.CloudVmRayResourceHandle,
+                    service_name: str) -> int:
+        assert handle.is_grpc_enabled_with_flag
+        request = servev1_pb2.AddVersionRequest(service_name=service_name)
+        response = backend_utils.invoke_skylet_with_retries(
+            lambda: backends.SkyletClient(handle.get_grpc_channel()
+                                         ).add_serve_version(request))
+        return response.version
+
+    @classmethod
+    def terminate_services(cls, handle: backends.CloudVmRayResourceHandle,
+                           service_names: Optional[List[str]], purge: bool,
+                           pool: bool) -> str:
+        assert handle.is_grpc_enabled_with_flag
+        request = TerminateServicesRequestConverter.to_proto(
+            service_names, purge, pool)
+        response = backend_utils.invoke_skylet_with_retries(
+            lambda: backends.SkyletClient(handle.get_grpc_channel()
+                                         ).terminate_services(request))
+        return response.message
+
+    @classmethod
+    def terminate_replica(cls, handle: backends.CloudVmRayResourceHandle,
+                          service_name: str, replica_id: int,
+                          purge: bool) -> str:
+        assert handle.is_grpc_enabled_with_flag
+        request = servev1_pb2.TerminateReplicaRequest(service_name=service_name,
+                                                      replica_id=replica_id,
+                                                      purge=purge)
+        response = backend_utils.invoke_skylet_with_retries(
+            lambda: backends.SkyletClient(handle.get_grpc_channel()
+                                         ).terminate_replica(request))
+        return response.message
+
+    @classmethod
+    def wait_service_registration(cls,
+                                  handle: backends.CloudVmRayResourceHandle,
+                                  service_name: str, job_id: int,
+                                  pool: bool) -> int:
+        assert handle.is_grpc_enabled_with_flag
+        request = servev1_pb2.WaitServiceRegistrationRequest(
+            service_name=service_name, job_id=job_id, pool=pool)
+        response = backend_utils.invoke_skylet_with_retries(
+            lambda: backends.SkyletClient(handle.get_grpc_channel()
+                                         ).wait_service_registration(request))
+        return response.lb_port
+
+    @classmethod
+    def update_service(cls, handle: backends.CloudVmRayResourceHandle,
+                       service_name: str, version: int,
+                       mode: serve_utils.UpdateMode, pool: bool) -> None:
+        assert handle.is_grpc_enabled_with_flag
+        request = servev1_pb2.UpdateServiceRequest(service_name=service_name,
+                                                   version=version,
+                                                   mode=mode.value,
+                                                   pool=pool)
+        backend_utils.invoke_skylet_with_retries(lambda: backends.SkyletClient(
+            handle.get_grpc_channel()).update_service(request))

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)