imb 1.0.0__py3-none-any.whl → 1.0.1__py3-none-any.whl

imb/__init__.py

@@ -1 +0,0 @@
-from .inference_clients import OnnxClient, TritonClient

imb/base.py

@@ -0,0 +1,129 @@
+from typing import Dict, List, Optional, Tuple, Union
+import numpy as np
+import time
+import os
+
+
+class BaseClient:
+    def __init__(self, *args, **kwargs):
+        self.show_fps: bool = os.environ.get('SHOW_FPS') in {'yes', 'true'}
+        self.model_name = ''
+        self.fixed_batch = False
+        self.max_batch_size = 1
+        self.is_async = False
+        self.return_dict = True
+
+        self.inputs_names: List[str] = []
+        self.inputs_shapes: List[tuple] = []
+        self.np_inputs_dtypes: List[np.dtype] = []
+
+        self.outputs_names: List[str] = []
+        self.outputs_shapes: List[tuple] = []
+        self.np_outputs_dtypes: List[np.dtype] = []
+
+        self.sample_inputs: Optional[List[np.ndarray]] = None
+    
+    def _load_model_params(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def _create_input_sample(self):
+        if self.sample_inputs is not None:
+            # validate sample inputs
+            for sample_array, config_input_shape in zip(self.sample_inputs, self.inputs_shapes):
+                for i, (s_dim, t_dim) in enumerate(zip(sample_array.shape, config_input_shape)):
+                    if i == 0:
+                        if self.fixed_batch:
+                            assert s_dim == t_dim, \
+                                f'model support fixed batch size {t_dim}, \
+                                    sample_inputs has batch size {s_dim}'
+                        else:
+                            assert s_dim <= t_dim, \
+                                f'model support max batch size {t_dim}, \
+                                    sample_inputs has batch size {s_dim}'
+                        continue
+                    assert ((t_dim != -1) and (int(s_dim) == int(t_dim))) or t_dim == -1, \
+                        f'incorrect shape in sample_inputs {sample_array.shape}, must be {config_input_shape}'
+        else:
+            has_dynamic_shapes = any(
+                -1 in config_input_shape for config_input_shape in self.inputs_shapes
+            )
+            if has_dynamic_shapes:
+                return
+            self.sample_inputs = []
+            for config_input_shape, np_input_format in zip(self.inputs_shapes, self.np_inputs_dtypes):
+                self.sample_inputs.append(
+                    np.ones(config_input_shape).astype(np_input_format)
+                )
+    
+    def _create_batches(self, *inputs_data: np.ndarray) -> Tuple[Dict[str, List[np.ndarray]], List[int]]:
+        inputs_batches = dict()
+        paddings = []
+        for input_data, np_format, input_name in zip(inputs_data, self.np_inputs_dtypes, self.inputs_names):
+            input_data = input_data.astype(np_format)
+            input_batches, input_paddings = self._split_on_batches(input_data)
+            if paddings == []:
+                paddings = input_paddings
+            inputs_batches[input_name] = input_batches
+        return inputs_batches, paddings
+
+    def log(self, text, warn=False, err=False):
+        text = f'Model ({self.model_name}) - {text}'
+        if err:
+            print('error', text)
+        elif warn:
+            print('warning',text)
+        else:
+            print('debug', text)
+
+    def warmup_model(self):
+        if self.sample_inputs is None:
+            print('Model was not warmed up, because sample_inputs didn\'t set or shape is dynamic and cannot auto generate')
+            return
+        exception = None
+        for _ in range(5):
+            try:
+                _ = self.__call__(*self.sample_inputs)
+                exception = None
+            except Exception as e:
+                print(f'{e} while warmup, repeat inference...')
+                exception = e
+                time.sleep(2)
+        if exception is not None:
+            raise exception
+    
+    def pad_batch(self, batch: np.ndarray):
+        padding_size = self.max_batch_size - batch.shape[0]
+        if padding_size > 0:
+            pad = np.zeros([padding_size, *batch.shape[1:]], dtype=batch.dtype)
+            batch = np.concatenate((batch, pad), axis=0)
+        return batch, padding_size
+    
+    def _split_on_batches(self, input_data: np.ndarray):
+        batches = []
+        paddings = []
+        for i in range(0, len(input_data), self.max_batch_size):
+            batch = input_data[i:i+self.max_batch_size]
+            batches.append(batch)
+            paddings.append(0)
+        
+        if self.fixed_batch:
+            batches[-1], paddings[-1] = self.pad_batch(batches[-1])
+
+        return batches, paddings
+    
+    def forward(self, *input_data):
+        raise NotImplementedError
+    
+    def async_forward(self, *input_data):
+        raise NotImplementedError
+
+    def __call__(self, *args, **kwargs) -> Union[Dict[str, np.ndarray], List[np.ndarray]]:
+        t1 = time.time()
+        forward_func = self.async_forward if self.is_async else self.forward
+        output = forward_func(*args, **kwargs)
+        if self.return_dict is False:
+            output = [output[output_name] for output_name in self.outputs_names]
+        t2 = time.time()
+        if self.show_fps:
+            self.log(f'Model: {self.model_name} fps {int(len(args[0])/(t2-t1))}')
+        return output

imb/inference_clients/__init__.py

@@ -1,2 +0,0 @@
-from .onnx import OnnxClient
-from .triton import TritonClient

imb/onnx.py

@@ -0,0 +1,99 @@
+from collections import defaultdict
+from typing import Dict, List
+from .base import BaseClient
+import onnxruntime as rt
+import numpy as np
+
+
+class OnnxClient(BaseClient):
+    def __init__(self, model_path: str, 
+                 model_name: str, 
+                 providers: List[str] = ['CUDAExecutionProvider', 'CPUExecutionProvider'], 
+                 max_batch_size: int = 1,
+                 return_dict: bool = True,
+                 fixed_batch: bool = False,
+                 warmup: bool = False
+                 ):
+        super().__init__()
+        self.model_name = model_name
+        self.model_path = model_path
+        self.providers = providers
+        self.return_dict = return_dict
+        self.max_batch_size = max_batch_size
+        self.fixed_batch = fixed_batch
+
+        self._load_model_params(max_batch_size)
+
+        self.sample_inputs = [np.zeros((*shape,), dtype=dtype) for shape, dtype in zip(self.inputs_shapes, self.np_inputs_dtypes)]
+
+        if warmup:
+            self.warmup_model()
+
+    def _load_model_params(self, max_batch_size: int = 1):
+        """
+        Load model parameters from onnx model
+
+        Args:
+            max_batch_size (int, optional): max batch size. Defaults to 1.
+
+        Raises:
+            ValueError: not support dynamic batch
+        """
+        sess_options = rt.SessionOptions()
+        self.onnx_model = rt.InferenceSession(
+            self.model_path,
+            providers=self.providers,
+            sess_options=sess_options
+        )
+
+        model_inputs = self.onnx_model.get_inputs()
+        data_dtype = np.float16 if 'float16' in model_inputs[0].type else np.float32
+        self.inputs_names = [model_inputs[i].name for i in range(len(model_inputs))]
+        self.np_inputs_dtypes = [data_dtype for _ in range(len(self.inputs_names))]
+        self.inputs_shapes = [model_inputs[i].shape for i in range(len(model_inputs))] 
+        for i_input, shape in enumerate(self.inputs_shapes):
+            new_shape = []
+            for i_dim, value in enumerate(shape):
+                if isinstance(value, int):
+                    if i_dim == 0:
+                        self.max_batch_size = value
+                        self.log(f'set batch size {value} from model metadata')
+                    new_shape.append(value)
+                elif isinstance(value, str) and 'batch' in value:
+                    new_shape.append(max_batch_size)
+                    self.log(f'set batch size {value} from user settings')
+                else:
+                    raise ValueError(f'not support value {value} in input shape {shape}')
+            self.inputs_shapes[i_input] = new_shape            
+
+        model_outputs = self.onnx_model.get_outputs()
+        self.outputs_names = [model_outputs[i].name for i in range(len(model_outputs))]
+        self.np_outputs_dtypes = [data_dtype for _ in range(len(self.outputs_names))]
+    
+    def forward(self, *inputs_data: np.ndarray) -> Dict[str, np.ndarray]:
+        inputs_batches, batches_paddings = self._create_batches(*inputs_data)
+
+        result = defaultdict(list)
+        count_batches = len(next(iter(inputs_batches.values())))
+
+        for i_batch in range(count_batches):
+            batch = dict()
+            for input_name, np_dtype in zip(self.inputs_names, self.np_inputs_dtypes):
+                batch[input_name] = inputs_batches[input_name][i_batch].astype(np_dtype)
+
+            batch_result = self.onnx_model.run(self.outputs_names, batch)
+            batch_result = {
+                self.outputs_names[i]: batch_result[i].astype(self.np_outputs_dtypes[i])
+                for i in range(len(self.outputs_names))
+            }
+            
+            padding_size = batches_paddings[i_batch]
+            for output_name, output_value in batch_result.items():
+                result[output_name].append(
+                    output_value if padding_size == 0 else output_value[:-padding_size]
+                    )
+
+        for output_name, output_values in result.items(): 
+            result[output_name] = np.concatenate(output_values)
+
+        return result

imb/triton.py

@@ -0,0 +1,460 @@
+from collections import defaultdict
+from typing import Any, Dict, List, Literal, Optional, Tuple
+import tritonclient.http as httpclient
+import tritonclient.grpc as grpcclient
+import tritonclient.utils.cuda_shared_memory as cudashm
+from google.protobuf.json_format import MessageToJson
+from tritonclient import utils
+from .base import BaseClient
+import numpy as np
+import json
+import time
+
+
+class TritonClient(BaseClient):
+    def __init__(self, url: str,
+                 model_name: str,
+                 max_batch_size: int = 0,
+                 sample_inputs: Optional[List[np.ndarray]] = None,
+                 timeout: int = 10,
+                 resend_count: int = 10,
+                 fixed_batch: bool = True,
+                 is_async: bool = False,
+                 cuda_shm: bool = False,
+                 max_shm_regions: int = 2,
+                 scheme: Literal["http", "grpc"] = "http",
+                 return_dict: bool = True,
+                 warmup: bool = False
+                 ):
+        super().__init__()
+        self.model_name = model_name
+        self.scheme = scheme
+        self.client_module = httpclient if scheme == "http" else grpcclient
+        self.url = url
+        self.is_async = is_async
+        self.cuda_shm = cuda_shm
+        self.triton_timeout = timeout
+        self.resend_count = resend_count
+        self.max_shm_regions = max_shm_regions
+        self.return_dict = return_dict
+
+        self.triton_client = None
+        self._init_triton()
+
+        self.triton_inputs_dtypes = None
+        self.np_inputs_dtypes = None
+        
+        self.inputs_shapes = None
+        self.fixed_batch = fixed_batch
+        
+        self.inputs_names = None
+        self.outputs_names = None
+        
+        self.sample_inputs = sample_inputs       
+        
+        self._load_model_params(max_batch_size) 
+        self._create_input_sample()
+        if warmup:
+            self.warmup_model()
+
+        self.input_shm_handles = [None for _ in range(len(self.inputs_names))]
+        self.output_shm_handles = [None for _ in range(len(self.outputs_names))]
+
+        if self.cuda_shm:
+            assert is_async == False and fixed_batch == True
+            self._fill_output_dynamic_axis()
+            self._create_input_output_shm_handles()
+            self._register_cuda_shm_regions()
+
+    def io_summary(self):
+        return {
+            "model_name": self.model_name,
+            "url": self.url,
+            "scheme": self.scheme,
+
+            "inputs_shapes": self.inputs_shapes,
+            "inputs_names": self.inputs_names,
+            "triton_inputs_dtypes": self.triton_inputs_dtypes,
+            "np_inputs_dtypes": self.np_inputs_dtypes,
+
+            "outputs_shapes": self.outputs_shapes,
+            "outputs_names": self.outputs_names,
+            "triton_outputs_dtypes": self.triton_outputs_dtypes,
+            "np_outputs_dtypes": self.np_outputs_dtypes,
+
+            "fixed_batch": self.fixed_batch,
+            "async": self.is_async,
+            "cuda_shm": self.cuda_shm,
+            "max_shm_regions": self.max_shm_regions,
+        }
+
+    def _init_triton(self):
+        if self.triton_client is not None:
+            # reinit
+            self.triton_client.close()
+            time.sleep(3)
+
+        self.triton_client = self.client_module.InferenceServerClient(
+                                url=self.url,
+                                verbose=False,
+                                ssl=False,
+                                network_timeout=self.triton_timeout,
+                                connection_timeout=self.triton_timeout
+                            )
+
+    def _load_model_params(self, user_max_batch_size: int) -> None:
+        """
+        Load the model config from Triton Inferernce Server and update the class attributes.
+
+        Args:
+            user_max_batch_size (int): max_batch_size defined by user
+        """
+        if self.scheme == "grpc":
+            config = self.triton_client.get_model_config(self.model_name, as_json=True)
+            config = config["config"]
+        else:
+            config = self.triton_client.get_model_config(self.model_name)
+        
+        self.triton_inputs_dtypes, self.np_inputs_dtypes, \
+            self.inputs_shapes, self.inputs_names \
+                = self._parse_io_params(config['input'])
+        
+        self.triton_outputs_dtypes, self.np_outputs_dtypes, \
+            self.outputs_shapes, self.outputs_names \
+                = self._parse_io_params(config['output'])
+        
+        not_support_dynamic_batch = config['max_batch_size'] == 0
+        if not_support_dynamic_batch:
+            # use batch size from config
+            self.max_batch_size = config['input'][0]['dims'][0]
+            self.fixed_batch = True
+        else:
+            # user can decrease max_batch_size from config
+            if user_max_batch_size > 0:
+                self.max_batch_size = min(config['max_batch_size'], user_max_batch_size)
+            else:
+                self.max_batch_size = config['max_batch_size']
+            # in config's shape has no batch size
+            self.inputs_shapes = self._insert_batch_size_to_shapes(
+                self.inputs_shapes, self.max_batch_size
+                )
+            self.outputs_shapes = self._insert_batch_size_to_shapes(
+                self.outputs_shapes, self.max_batch_size
+                )
+    
+    def _fill_output_dynamic_axis(self) -> None:
+        """
+        Fill real values in the dynamic axis of the output shapes.
+        """
+        has_dynamic_shapes = any(
+                -1 in output_shape for output_shape in self.outputs_shapes
+            )
+        if has_dynamic_shapes:
+            start_cuda_shm_flag = self.cuda_shm
+            self.cuda_shm = False
+            outputs = self.forward(*self.sample_inputs)
+            self.outputs_shapes = [
+                list(outputs[output_name].shape) for output_name in self.outputs_names
+                ]
+            self.cuda_shm = start_cuda_shm_flag
+
+    @staticmethod
+    def _parse_io_params(io_params: List[Dict]) -> Tuple[List[str], List[np.dtype], List[List[int]], List[str]]:
+        """
+        Parse the input/output parameters from the model config.
+
+        Args:
+            io_params (List[Dict]): The input/output parameters.
+
+        Returns:
+            Tuple[List[str], List[np.dtype], List[List[int]], List[str]]: The input/output dtypes, shapes, and names.
+        """
+        triton_dtypes = []
+        np_dtypes = []
+        shapes = []
+        names = []
+        for params in io_params:
+            triton_dtypes.append(params['data_type'].replace('TYPE_', ''))
+            np_dtypes.append(utils.triton_to_np_dtype(triton_dtypes[-1]))
+            shapes.append(params['dims'])
+            names.append(params['name'])
+
+        return triton_dtypes, np_dtypes, shapes, names
+
+    @staticmethod
+    def _insert_batch_size_to_shapes(shapes: List[List], insert_batch_size: int) -> List[List[int]]:
+        """
+        Insert the batch size to the shapes.
+
+        Args:
+            shapes (List[List]): Shapes from config
+            insert_batch_size (int): Value for insert batch size to shape
+
+        Returns:
+            List[List[int]]: Result shape
+        """
+        return [[insert_batch_size] + shape for shape in shapes]
+
+    def _generate_shm_name(self, ioname: str) -> str:
+        """
+        Generate shared region name
+
+        Args:
+            ioname (str): Input/output name
+
+        Returns:
+            str: Shared region name
+        """
+        return f'{self.model_name}_{ioname}_{time.time()}'
+
+    def _get_old_regions_names(self, regions_statuses: list, new_triton_shm_name: str) -> List[str]:
+        """
+        Get old regions names for unregister
+
+        Args:
+            regions_statuses (list): responce of get_cuda_shared_memory_status from triton
+            new_triton_shm_name (str): name of new region
+
+        Returns:
+            List[str]: old regions names for unregister
+        """
+        i_sep = len(new_triton_shm_name) - new_triton_shm_name[::-1].index('_') - 1
+        region_name = new_triton_shm_name[:i_sep]
+        registrated_regions = [
+            (region['name'], float(region['name'][i_sep+1:])) 
+            for region in regions_statuses if region['name'].startswith(region_name)
+        ]
+        registrated_regions.sort(key=lambda x: x[1])
+        count_old_regions = len(registrated_regions) - self.max_shm_regions + 1
+        old_regions = []
+        if count_old_regions > 0:
+            old_regions = [name for name, _ in registrated_regions[:count_old_regions]]
+        return old_regions
+
+    def _register_cuda_shm_regions(self):
+        """
+        Register CUDA shared memory regions in Triton
+        """
+        if self.scheme == "grpc":
+            regions_statuses = self.triton_client.get_cuda_shared_memory_status(as_json=True)['regions']
+        else:
+            regions_statuses = self.triton_client.get_cuda_shared_memory_status()
+
+        for shm_handle in self.input_shm_handles + self.output_shm_handles:
+            old_regions_names = self._get_old_regions_names(regions_statuses, shm_handle._triton_shm_name)
+            for old_region_name in old_regions_names:
+                self.triton_client.unregister_cuda_shared_memory(old_region_name)
+            self.triton_client.register_cuda_shared_memory(
+                shm_handle._triton_shm_name, cudashm.get_raw_handle(shm_handle), 0, shm_handle._byte_size
+            )
+
+    def _create_cuda_shm_handle(self, shape: List[int], dtype: np.dtype, name: str) -> Any:
+        """
+        Create CUDA shared memory handle
+
+        Args:
+            shape (List[int]): Shape of cuda shared memory region
+            dtype (np.dtype): Data type of input/output data
+            name (str): Input/output name
+
+        Returns:
+            Any: CUDA shared memory handle
+        """
+        byte_size = int(np.prod(shape) * np.dtype(dtype).itemsize)
+        shm_name = self._generate_shm_name(name)
+        return cudashm.create_shared_memory_region(shm_name, byte_size, 0)
+
+    def _create_cuda_shm_handles_for_io(self, shapes: List[List[int]], 
+                                        dtypes: List[np.dtype], 
+                                        names: List[str]) -> List[Any]:
+        """
+        Create CUDA shared memory handles for inputs or outputs
+
+        Args:
+            shapes (List[List[int]]): Shapes of cuda shared memory regions
+            dtypes (List[np.dtype]): Data types of input/output data
+            names (List[str]): Input/output names
+
+        Returns:
+            List[Any]: CUDA shared memory handles
+        """
+        return [self._create_cuda_shm_handle(shape, dtype, name) 
+                for shape, dtype, name in zip(shapes, dtypes, names)]
+
+    def _create_input_output_shm_handles(self) -> None:
+        """
+        Create CUDA shared memory handles for inputs and outputs
+        """
+        self.input_shm_handles = self._create_cuda_shm_handles_for_io(
+            self.inputs_shapes, self.np_inputs_dtypes, self.inputs_names
+        )
+        self.output_shm_handles = self._create_cuda_shm_handles_for_io(
+            self.outputs_shapes, self.np_outputs_dtypes, self.outputs_names
+        )
+
+    def _create_triton_input(self, input_data: np.ndarray, input_name: str, 
+                             config_input_format: str, shm_handle = None) -> Any:
+        """
+        Create triton InferInput
+
+        Args:
+            input_data (np.ndarray): data for send to model
+            input_name (str): name of input
+            config_input_format (str): triton input format
+            shm_handle (_type_, optional): CUDA shared memory handle. Defaults to None.
+
+        Returns:
+            Any: triton InferInput for sending request
+        """
+        infer_input = self.client_module.InferInput(input_name, input_data.shape, config_input_format)
+        if self.cuda_shm:
+            cudashm.set_shared_memory_region(shm_handle, [input_data])
+            infer_input.set_shared_memory(shm_handle._triton_shm_name, shm_handle._byte_size)
+        else:
+            infer_input.set_data_from_numpy(input_data)
+        return infer_input
+
+    def _create_triton_output(self, output_name: str, binary: bool = True, shm_handle = None) -> Any:
+        """
+        Create triton InferRequestedOutput
+
+        Args:
+            output_name (str): output name
+            binary (bool, optional): Whether the output is binary. Defaults to True.
+            shm_handle (_type_, optional): CUDA shared memory handle. Defaults to None.
+
+        Returns:
+            Any: triton InferRequestedOutput for receiving response
+        """
+        if self.scheme == "grpc":
+            infer_output = self.client_module.InferRequestedOutput(output_name)
+        else:
+            infer_output = self.client_module.InferRequestedOutput(output_name, binary_data=binary)
+        if self.cuda_shm:
+            infer_output.set_shared_memory(shm_handle._triton_shm_name, shm_handle._byte_size)
+        return infer_output
+
+    def _postprocess_triton_result(self, triton_response: Any, padding_size: int) -> Dict[str, np.ndarray]:
+        """
+        Postprocess triton response.
+
+        Args:
+            triton_response (Any): triton response
+            padding_size (int): padding size for unpad output data
+
+        Returns:
+            Dict[str, np.ndarray]: dict of output name and output data
+        """
+        result = dict()
+        for output_name, shm_op_handle in zip(self.outputs_names, self.output_shm_handles):
+            if self.cuda_shm:
+                if self.scheme == "grpc":
+                    # output = triton_response.get_output(output_name, as_json=True) # WARN: bug in tritonclient library, return None
+                    output = json.loads(MessageToJson(triton_response.get_output(output_name)))
+                else:
+                    output = triton_response.get_output(output_name)
+                result[output_name] = cudashm.get_contents_as_numpy(
+                    shm_op_handle,
+                    utils.triton_to_np_dtype(output["datatype"]),
+                    output["shape"],
+                )
+            else:
+                result[output_name] = triton_response.as_numpy(output_name)
+
+            if padding_size != 0:
+                result[output_name] = result[output_name][:-padding_size]
+                
+        return result
+
+    def forward(self, *inputs_data: np.ndarray) -> Dict[str, np.ndarray]:
+        assert len(inputs_data) == len(self.inputs_names), 'inputs number is not equal to model inputs'
+        inputs_batches, batches_paddings = self._create_batches(*inputs_data)
+
+        result = defaultdict(list)
+        count_batches = len(next(iter(inputs_batches.values())))
+        
+        for i_batch in range(count_batches):
+            triton_inputs = []
+            for input_name, config_input_format, shm_ip_handle in \
+                    zip(self.inputs_names, self.triton_inputs_dtypes, self.input_shm_handles):
+                triton_input = self._create_triton_input(
+                    inputs_batches[input_name][i_batch], input_name, config_input_format, shm_ip_handle
+                    )
+                triton_inputs.append(triton_input)
+
+            triton_outputs = []
+            for output_name, shm_op_handle in zip(self.outputs_names, self.output_shm_handles):
+                triton_output = self._create_triton_output(
+                    output_name, binary=True, shm_handle=shm_op_handle
+                    )
+                triton_outputs.append(triton_output)
+
+            triton_response = self.triton_client.infer(
+                model_name=self.model_name, 
+                inputs=triton_inputs, 
+                outputs=triton_outputs
+                )
+            
+            batch_result = self._postprocess_triton_result(triton_response, batches_paddings[i_batch])
+
+            for output_name, output_value in batch_result.items():
+                result[output_name].append(output_value)
+
+        for output_name, output_values in result.items(): 
+            result[output_name] = np.concatenate(output_values)
+
+        return result
+
+    def send_async_requests(self, inputs_batches: Dict):
+        count_batches = len(next(iter(inputs_batches.values())))
+        
+        triton_response_handles = []
+
+        for i_batch in range(count_batches):
+            triton_inputs = []
+            for input_name, config_input_format, shm_ip_handle in \
+                    zip(self.inputs_names, self.triton_inputs_dtypes, self.input_shm_handles):
+                triton_input = self._create_triton_input(
+                    inputs_batches[input_name][i_batch], input_name, config_input_format, shm_ip_handle
+                    )
+                triton_inputs.append(triton_input)
+
+            triton_outputs = []
+            for output_name, shm_op_handle in zip(self.outputs_names, self.output_shm_handles):
+                triton_output = self._create_triton_output(
+                    output_name, binary=True, shm_handle=shm_op_handle
+                    )
+                triton_outputs.append(triton_output)
+            
+            triton_response_handle = self.triton_client.async_infer(
+                        model_name=self.model_name, 
+                        inputs=triton_inputs, 
+                        outputs=triton_outputs
+                        )
+            triton_response_handles.append(triton_response_handle)
+        
+        return triton_response_handles
+    
+    def get_async_results(self, triton_response_handles, batches_paddings):
+        result = defaultdict(list)
+        for i_batch, triton_response_handle in enumerate(triton_response_handles):
+            triton_response = triton_response_handle.get_result()
+            batch_result = self._postprocess_triton_result(triton_response, batches_paddings[i_batch])
+
+            for output_name, output_value in batch_result.items():
+                result[output_name].append(output_value)
+
+        for output_name, output_values in result.items(): 
+            result[output_name] = np.concatenate(output_values)
+        
+        return result
+    
+    def async_forward(self, *inputs_data: np.ndarray):
+        assert len(inputs_data) == len(self.inputs_names), 'inputs number is not equal to model inputs'
+        inputs_batches, batches_paddings = self._create_batches(*inputs_data)
+        
+        triton_response_handles = self.send_async_requests(inputs_batches)
+
+        result = self.get_async_results(triton_response_handles, batches_paddings)
+
+        return result

imb-1.0.1.dist-info/METADATA

@@ -0,0 +1,105 @@
+Metadata-Version: 2.2
+Name: imb
+Version: 1.0.1
+Summary: Python library for run inference of deep learning models in different backends
+Home-page: https://github.com/TheConstant3/InferenceMultiBackend
+Author: p-constant
+Author-email: nikshorop@gmail.com
+Classifier: Programming Language :: Python :: 3.8
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy
+Provides-Extra: triton
+Requires-Dist: tritonclient[all]>=2.38.0; extra == "triton"
+Provides-Extra: onnxcpu
+Requires-Dist: onnxruntime>=1.16.0; extra == "onnxcpu"
+Provides-Extra: onnxgpu
+Requires-Dist: onnxruntime-gpu>=1.16.0; extra == "onnxgpu"
+Provides-Extra: all
+Requires-Dist: tritonclient[all]>=2.38.0; extra == "all"
+Requires-Dist: onnxruntime>=1.16.0; extra == "all"
+Requires-Dist: onnxruntime-gpu>=1.16.0; extra == "all"
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: provides-extra
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# InferenceMultiBackend
+
+Python library for run inference of deep learning models in different backends
+
+## Installation
+
+For use triton inference client:
+```pip install imb[triton]```
+
+For use onnxruntime-gpu client:
+```pip install imb[onnxgpu]```
+
+For use onnxruntime client:
+```pip install imb[onnxcpu]```
+
+For support all implemented clients:
+```pip install imb[all]```
+
+## Usage
+
+OnnxClient usage example
+```
+onnx_client = OnnxClient(
+    model_path='model.onnx',
+    model_name='any name',
+    providers=['CUDAExecutionProvider', 'CPUExecutionProvider'],
+    max_batch_size=16,
+    return_dict=True,
+    fixed_batch=True,
+    warmup=True
+)
+# if model has fixed input size (except batch size) then sample_inputs will be created
+sample_inputs = onnx_client.sample_inputs
+print('inputs shapes', [o.shape for o in sample_inputs])
+outputs = onnx_client(*sample_inputs)
+print('outputs shapes', [(o_name, o_value.shape) for o_name, o_value in outputs.items()])
+```
+
+TritonClient usage example
+```
+triton_client = TritonClient(
+    url='localhost:8000',
+    model_name='arcface',
+    max_batch_size=16,
+    timeout=10,
+    resend_count=10,
+    fixed_batch=True,
+    is_async=False,
+    cuda_shm=False,
+    max_shm_regions=2,
+    scheme='http',
+    return_dict=True,
+    warmup=False
+)
+# if model has fixed input size (except batch size) then sample_inputs will be created
+sample_inputs = triton_client.sample_inputs
+print('inputs shapes', [o.shape for o in sample_inputs])
+outputs = triton_client(*sample_inputs)
+print('outputs shapes', [(o_name, o_value.shape) for o_name, o_value in outputs.items()])
+```
+
+## Notes
+
+max_batch_size - maximum batch size for inference. If input data larger that max_batch_size, then input data will be splitted to several batches.
+
+fixed_batch - if fixed batch is True, then each batch will have fixed size (padding the smallest batch to max_batch_size).
+
+warmup - if True, model will run several calls on sample_inputs while initialization. 
+
+return_dict - if True, __call__ return dict {'output_name1': output_value1, ...}, else [output_value1, ...]

imb-1.0.1.dist-info/RECORD

@@ -0,0 +1,13 @@
+imb/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+imb/base.py,sha256=oBmiTu4rHgzED5kCxKPvS9e3PhI229Pj5lxuPm7ep6M,5189
+imb/onnx.py,sha256=g3vQBJPeln0YUOQ1X9RjZce8AAi-7SXntLyevOZZdG8,4100
+imb/triton.py,sha256=hdnCtDjoRAl_Ss49_ayvW3-VhsYcY2MbNqh3ax6y-18,18629
+imb/inference_clients/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+imb/inference_clients/base.py,sha256=oBmiTu4rHgzED5kCxKPvS9e3PhI229Pj5lxuPm7ep6M,5189
+imb/inference_clients/onnx.py,sha256=g3vQBJPeln0YUOQ1X9RjZce8AAi-7SXntLyevOZZdG8,4100
+imb/inference_clients/triton.py,sha256=hdnCtDjoRAl_Ss49_ayvW3-VhsYcY2MbNqh3ax6y-18,18629
+imb-1.0.1.dist-info/LICENSE,sha256=pAZXnNE2dxxwXFIduGyn1gpvPefJtUYOYZOi3yeGG94,1068
+imb-1.0.1.dist-info/METADATA,sha256=5x0Xa-Gbg8D8e0mNMzKZB4l27xwNlKfyUXzDwQfkJUA,3240
+imb-1.0.1.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+imb-1.0.1.dist-info/top_level.txt,sha256=kY8Fp1i_MzTZhuoVhVexG762D8HBd-THfX_lfw4EZmY,4
+imb-1.0.1.dist-info/RECORD,,

imb-1.0.0.dist-info/METADATA

@@ -1,30 +0,0 @@
-Metadata-Version: 2.2
-Name: imb
-Version: 1.0.0
-Summary: Python library for run inference of deep learning models in different backends
-Home-page: https://github.com/TheConstant3/InferenceMultiBackend
-Author: p-constant
-Author-email: nikshorop@gmail.com
-Classifier: Programming Language :: Python :: 3.8
-Classifier: License :: OSI Approved :: MIT License
-Classifier: Operating System :: OS Independent
-Requires-Python: >=3.8
-Description-Content-Type: text/markdown
-License-File: LICENSE
-Requires-Dist: onnxruntime-gpu>=1.16.0
-Requires-Dist: tritonclient[all]>=2.38.0
-Requires-Dist: numpy>=1.19.4
-Dynamic: author
-Dynamic: author-email
-Dynamic: classifier
-Dynamic: description
-Dynamic: description-content-type
-Dynamic: home-page
-Dynamic: requires-dist
-Dynamic: requires-python
-Dynamic: summary
-
-# InferenceMultiBackend
-
-Python library for run inference of deep learning models in different backends
-

imb-1.0.0.dist-info/RECORD

@@ -1,10 +0,0 @@
-imb/__init__.py,sha256=8XoaonMp09UWmynubLMIu2bln41iKgIdWj-wxgsQjnk,55
-imb/inference_clients/__init__.py,sha256=Glv4yD0QdtZmCOiYFbILSl90VhxdwvPoH9gFczHlVFk,61
-imb/inference_clients/base.py,sha256=oBmiTu4rHgzED5kCxKPvS9e3PhI229Pj5lxuPm7ep6M,5189
-imb/inference_clients/onnx.py,sha256=g3vQBJPeln0YUOQ1X9RjZce8AAi-7SXntLyevOZZdG8,4100
-imb/inference_clients/triton.py,sha256=hdnCtDjoRAl_Ss49_ayvW3-VhsYcY2MbNqh3ax6y-18,18629
-imb-1.0.0.dist-info/LICENSE,sha256=pAZXnNE2dxxwXFIduGyn1gpvPefJtUYOYZOi3yeGG94,1068
-imb-1.0.0.dist-info/METADATA,sha256=NZcJPx91mzPg4Zo9FZxlMQE4c6zB2s_yPVhhRVxPBzM,898
-imb-1.0.0.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-imb-1.0.0.dist-info/top_level.txt,sha256=kY8Fp1i_MzTZhuoVhVexG762D8HBd-THfX_lfw4EZmY,4
-imb-1.0.0.dist-info/RECORD,,