# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
"""Client and server classes corresponding to protobufs-defined services."""
# Third Party
import grpc

# Local
from . import model_runtime_pb2 as model__runtime__pb2


class ModelRuntimeStub(object):
    """this is the internal "sidecar" API for interfacing with a
    colocated model runtime container
    """

    def __init__(self, channel):
        """Constructor.

        Args:
            channel: A grpc.Channel.
        """
        self.loadModel = channel.unary_unary(
            "/mmesh.ModelRuntime/loadModel",
            request_serializer=model__runtime__pb2.LoadModelRequest.SerializeToString,
            response_deserializer=model__runtime__pb2.LoadModelResponse.FromString,
        )
        self.unloadModel = channel.unary_unary(
            "/mmesh.ModelRuntime/unloadModel",
            request_serializer=model__runtime__pb2.UnloadModelRequest.SerializeToString,
            response_deserializer=model__runtime__pb2.UnloadModelResponse.FromString,
        )
        self.predictModelSize = channel.unary_unary(
            "/mmesh.ModelRuntime/predictModelSize",
            request_serializer=model__runtime__pb2.PredictModelSizeRequest.SerializeToString,
            response_deserializer=model__runtime__pb2.PredictModelSizeResponse.FromString,
        )
        self.modelSize = channel.unary_unary(
            "/mmesh.ModelRuntime/modelSize",
            request_serializer=model__runtime__pb2.ModelSizeRequest.SerializeToString,
            response_deserializer=model__runtime__pb2.ModelSizeResponse.FromString,
        )
        self.runtimeStatus = channel.unary_unary(
            "/mmesh.ModelRuntime/runtimeStatus",
            request_serializer=model__runtime__pb2.RuntimeStatusRequest.SerializeToString,
            response_deserializer=model__runtime__pb2.RuntimeStatusResponse.FromString,
        )


class ModelRuntimeServicer(object):
    """this is the internal "sidecar" API for interfacing with a
    colocated model runtime container
    """

    def loadModel(self, request, context):
        """Load a model, return when model is fully loaded.
        Include size of loaded model in response if no additional cost.
        A gRPC error code of PRECONDITION_FAILED or INVALID_ARGUMENT
        should be returned if no attempt to load the model was made
        (so can be sure that no space remains used).
        Note that the RPC may be cancelled by model-mesh prior to completion,
        after which an unloadModel call will immediately be sent for the same model.
        To avoid state inconsistency and "leaking" memory, implementors should
        ensure that this case is properly handled, i.e. that the model doesn't
        remain loaded after returning successfully from this unloadModel call.
        """
        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
        context.set_details("Method not implemented!")
        raise NotImplementedError("Method not implemented!")

    def unloadModel(self, request, context):
        """Unload a previously loaded (or failed) model. Return when model
        is fully unloaded, or immediately if not found/loaded.
        """
        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
        context.set_details("Method not implemented!")
        raise NotImplementedError("Method not implemented!")

    def predictModelSize(self, request, context):
        """Predict size of not-yet-loaded model - must return almost immediately.
        Should not perform expensive computation or remote lookups.
        Should be a conservative estimate.
        """
        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
        context.set_details("Method not implemented!")
        raise NotImplementedError("Method not implemented!")

    def modelSize(self, request, context):
        """Calculate size (memory consumption) of currently-loaded model"""
        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
        context.set_details("Method not implemented!")
        raise NotImplementedError("Method not implemented!")

    def runtimeStatus(self, request, context):
        """Provide basic runtime status and parameters; called only during startup.
        Before returning a READY status, implementations should check for and
        purge any/all currently-loaded models. Since this is only called during
        startup, there should very rarely be any, but if there are it implies
        the model-mesh container restarted unexpectedly and such a purge must
        be done to ensure continued consistency of state and avoid over-committing
        resources.
        """
        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
        context.set_details("Method not implemented!")
        raise NotImplementedError("Method not implemented!")


def add_ModelRuntimeServicer_to_server(servicer, server):
    rpc_method_handlers = {
        "loadModel": grpc.unary_unary_rpc_method_handler(
            servicer.loadModel,
            request_deserializer=model__runtime__pb2.LoadModelRequest.FromString,
            response_serializer=model__runtime__pb2.LoadModelResponse.SerializeToString,
        ),
        "unloadModel": grpc.unary_unary_rpc_method_handler(
            servicer.unloadModel,
            request_deserializer=model__runtime__pb2.UnloadModelRequest.FromString,
            response_serializer=model__runtime__pb2.UnloadModelResponse.SerializeToString,
        ),
        "predictModelSize": grpc.unary_unary_rpc_method_handler(
            servicer.predictModelSize,
            request_deserializer=model__runtime__pb2.PredictModelSizeRequest.FromString,
            response_serializer=model__runtime__pb2.PredictModelSizeResponse.SerializeToString,
        ),
        "modelSize": grpc.unary_unary_rpc_method_handler(
            servicer.modelSize,
            request_deserializer=model__runtime__pb2.ModelSizeRequest.FromString,
            response_serializer=model__runtime__pb2.ModelSizeResponse.SerializeToString,
        ),
        "runtimeStatus": grpc.unary_unary_rpc_method_handler(
            servicer.runtimeStatus,
            request_deserializer=model__runtime__pb2.RuntimeStatusRequest.FromString,
            response_serializer=model__runtime__pb2.RuntimeStatusResponse.SerializeToString,
        ),
    }
    generic_handler = grpc.method_handlers_generic_handler(
        "mmesh.ModelRuntime", rpc_method_handlers
    )
    server.add_generic_rpc_handlers((generic_handler,))


# This class is part of an EXPERIMENTAL API.
class ModelRuntime(object):
    """this is the internal "sidecar" API for interfacing with a
    colocated model runtime container
    """

    @staticmethod
    def loadModel(
        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,
            "/mmesh.ModelRuntime/loadModel",
            model__runtime__pb2.LoadModelRequest.SerializeToString,
            model__runtime__pb2.LoadModelResponse.FromString,
            options,
            channel_credentials,
            insecure,
            call_credentials,
            compression,
            wait_for_ready,
            timeout,
            metadata,
        )

    @staticmethod
    def unloadModel(
        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,
            "/mmesh.ModelRuntime/unloadModel",
            model__runtime__pb2.UnloadModelRequest.SerializeToString,
            model__runtime__pb2.UnloadModelResponse.FromString,
            options,
            channel_credentials,
            insecure,
            call_credentials,
            compression,
            wait_for_ready,
            timeout,
            metadata,
        )

    @staticmethod
    def predictModelSize(
        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,
            "/mmesh.ModelRuntime/predictModelSize",
            model__runtime__pb2.PredictModelSizeRequest.SerializeToString,
            model__runtime__pb2.PredictModelSizeResponse.FromString,
            options,
            channel_credentials,
            insecure,
            call_credentials,
            compression,
            wait_for_ready,
            timeout,
            metadata,
        )

    @staticmethod
    def modelSize(
        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,
            "/mmesh.ModelRuntime/modelSize",
            model__runtime__pb2.ModelSizeRequest.SerializeToString,
            model__runtime__pb2.ModelSizeResponse.FromString,
            options,
            channel_credentials,
            insecure,
            call_credentials,
            compression,
            wait_for_ready,
            timeout,
            metadata,
        )

    @staticmethod
    def runtimeStatus(
        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,
            "/mmesh.ModelRuntime/runtimeStatus",
            model__runtime__pb2.RuntimeStatusRequest.SerializeToString,
            model__runtime__pb2.RuntimeStatusResponse.FromString,
            options,
            channel_credentials,
            insecure,
            call_credentials,
            compression,
            wait_for_ready,
            timeout,
            metadata,
        )
