# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/02_grpc.ipynb.

# %% auto 0
__all__ = ['TorchServeClientGRPC']

# %% ../nbs/02_grpc.ipynb 3
import os
import grpc
import torchserve_client.proto.inference_pb2 as inference_pb2
import torchserve_client.proto.inference_pb2_grpc as inference_pb2_grpc
import torchserve_client.proto.management_pb2 as management_pb2
import torchserve_client.proto.management_pb2_grpc as management_pb2_grpc

# %% ../nbs/02_grpc.ipynb 4
class BaseClient:
    def __init__(self, base_url=None):
        base_url = (
            base_url
            if base_url
            else os.environ.get("TORCHSERVE_URL", "http://localhost")
        )
        self.base_url = base_url.split("//")[1]

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}(base_url={self.base_url})"

    def _filter_none_values(self, data):
        return {key: value for key, value in data.items() if value is not None}

# %% ../nbs/02_grpc.ipynb 5
class ManagementClient(BaseClient):
    def __init__(self, base_url=None, port=7071):
        super().__init__(base_url)
        self.port = port
        self.base_url = f"{self.base_url}:{self.port}"
        self.channel = grpc.insecure_channel(self.base_url)
        self.stub = management_pb2_grpc.ManagementAPIsServiceStub(self.channel)

    def describe_model(self, model_name, model_version=None, customized=None):
        method_args = {key: value for key, value in locals().items() if key != "self"}
        method_args = self._filter_none_values(method_args)
        request = management_pb2.DescribeModelRequest(model_name=model_name)
        return self.stub.DescribeModel(request)

    def list_models(self, limit=None, next_page_token=None):
        method_args = {key: value for key, value in locals().items() if key != "self"}
        method_args = self._filter_none_values(method_args)
        request = management_pb2.ListModelsRequest(**method_args)
        return self.stub.ListModels(request)

    def register_model(
        self,
        batch_size=None,
        handler=None,
        initial_workers=None,
        max_batch_delay=None,
        model_name=None,
        response_timeout=None,
        runtime=None,
        synchronous=None,
        url=None,
        s3_sse_kms=None,
    ):
        method_args = {key: value for key, value in locals().items() if key != "self"}
        method_args = self._filter_none_values(method_args)
        request = management_pb2.RegisterModelRequest(**method_args)
        return self.stub.RegisterModel(request)

    def scale_worker(
        self,
        model_name,
        min_worker=None,
        max_worker=None,
        model_version=None,
        number_gpu=None,
        synchronous=None,
        timeout=None,
    ):
        method_args = {key: value for key, value in locals().items() if key != "self"}
        method_args = self._filter_none_values(method_args)
        request = management_pb2.ScaleWorkerRequest(**method_args)
        return self.stub.ScaleWorker(request)

    def set_default(self, model_name, model_version):
        method_args = {key: value for key, value in locals().items() if key != "self"}
        method_args = self._filter_none_values(method_args)
        request = management_pb2.SetDefaultRequest(**method_args)
        return self.stub.SetDefault(request)

    def unregister_model(self, model_name, model_version=None):
        method_args = {key: value for key, value in locals().items() if key != "self"}
        method_args = self._filter_none_values(method_args)
        request = management_pb2.UnregisterModelRequest(**method_args)
        return self.stub.UnregisterModel(request)

# %% ../nbs/02_grpc.ipynb 6
class InferenceClient(BaseClient):
    def __init__(self, base_url=None, port=7070):
        super().__init__(base_url)
        self.port = port
        self.base_url = f"{self.base_url}:{self.port}"
        self.channel = grpc.insecure_channel(self.base_url)
        self.stub = inference_pb2_grpc.InferenceAPIsServiceStub(self.channel)

    def ping(self):
        return self.stub.Ping()

    def predictions(self, model_name, input_data, model_version=None):
        """
        input_data = {"data": data}
        """
        method_args = {key: value for key, value in locals().items() if key != "self"}
        method_args = self._filter_none_values(method_args)
        request = inference_pb2.PredictionsRequest(**method_args)
        return self.stub.Predictions(request)

    def stream_predictions(self, model_name, input_data, model_version=None):
        """
        input_data = {"data": data}
        """
        method_args = {key: value for key, value in locals().items() if key != "self"}
        method_args = self._filter_none_values(method_args)
        request = inference_pb2.PredictionsRequest(**method_args)
        return self.stub.StreamPredictions(request)

# %% ../nbs/02_grpc.ipynb 7
class TorchServeClientGRPC:
    def __init__(self, base_url=None, management_port=7071, inference_port=7070):
        self.management = ManagementClient(base_url, management_port)
        self.inference = InferenceClient(base_url, inference_port)

    def __repr__(self):
        url = self.management.base_url.rsplit(":", 1)[0]
        return f"TorchServeClientGRPC(base_url={url}, management_port={self.management.port}, inference_port={self.inference.port})"
