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

imb/triton.py

@@ -3,6 +3,7 @@ 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
+import tritonclient.utils.shared_memory as shm
 from google.protobuf.json_format import MessageToJson
 from tritonclient import utils
 from .base import BaseClient
@@ -11,6 +12,14 @@ import json
 import time
 
 
+class ShmHandlerWrapper:
+    def __init__(self, handler: Any, name: str, size: int):
+        self.handler = handler
+        self.name = name
+        self.size = size
+
+
+
 class TritonClient(BaseClient):
     def __init__(self, url: str,
                  model_name: str,
@@ -20,19 +29,40 @@ class TritonClient(BaseClient):
                  resend_count: int = 10,
                  fixed_batch: bool = True,
                  is_async: bool = False,
-                 cuda_shm: bool = False,
-                 max_shm_regions: int = 2,
+                 use_cuda_shm: bool = False,
+                 use_system_shm: bool = False,
+                 max_shm_regions: int = 0,
                  scheme: Literal["http", "grpc"] = "http",
                  return_dict: bool = True,
                  warmup: bool = False
                  ):
+        """_summary_
+
+        Args:
+            url (str): url of the triton server
+            model_name (str): name of the model endpoint
+            max_batch_size (int, optional): max batch size. Defaults to 0 (get value from triton config).
+            sample_inputs (Optional[List[np.ndarray]], optional): samples for warmup. Defaults to None (zeros array).
+            timeout (int, optional): triton client timeout. Defaults to 10.
+            resend_count (int, optional): triton client resend count. Defaults to 10.
+            fixed_batch (bool, optional): use fixed batch size, using padding for smaller batch. Defaults to True.
+            is_async (bool, optional): async inference. Defaults to False.
+            use_cuda_shm (bool, optional): use cuda shared memory. Defaults to False.
+            use_system_shm (bool, optional): use system shared memory. Defaults to False.
+            max_shm_regions (int, optional): max clients for shared memory. Will unregister old regions. Defaults to 0.
+            scheme (Literal["http", "grpc"], optional): scheme for triton client. Defaults to "http".
+            return_dict (bool, optional): return dict or list of values. Defaults to True.
+            warmup (bool, optional): warmup model. Defaults to False.
+        """
         super().__init__()
+        assert not (use_cuda_shm and use_system_shm), 'shm and cuda_shm are mutually exclusive'
         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.use_cuda_shm = use_cuda_shm
+        self.use_system_shm = use_system_shm
         self.triton_timeout = timeout
         self.resend_count = resend_count
         self.max_shm_regions = max_shm_regions
@@ -57,14 +87,16 @@ class TritonClient(BaseClient):
         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))]
+        self.input_shm_handlers: List[Optional[ShmHandlerWrapper]] = \
+            [None for _ in range(len(self.inputs_names))]
+        self.output_shm_handlers: List[Optional[ShmHandlerWrapper]] = \
+            [None for _ in range(len(self.outputs_names))]
 
-        if self.cuda_shm:
+        if self.use_cuda_shm or self.use_system_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()
+            self._register_shm_regions()
 
     def io_summary(self):
         return {
@@ -84,7 +116,8 @@ class TritonClient(BaseClient):
 
             "fixed_batch": self.fixed_batch,
             "async": self.is_async,
-            "cuda_shm": self.cuda_shm,
+            "cuda_shm": self.use_cuda_shm,
+            "shm": self.use_system_shm,
             "max_shm_regions": self.max_shm_regions,
         }
 
@@ -150,13 +183,16 @@ class TritonClient(BaseClient):
                 -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
+            start_cuda_shm_flag = self.use_cuda_shm
+            start_system_shm_flag = self.use_system_shm
+            self.use_cuda_shm = False
+            self.use_system_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
+            self.use_cuda_shm = start_cuda_shm_flag
+            self.use_system_shm = start_system_shm_flag
 
     @staticmethod
     def _parse_io_params(io_params: List[Dict]) -> Tuple[List[str], List[np.dtype], List[List[int]], List[str]]:
@@ -212,12 +248,14 @@ class TritonClient(BaseClient):
         Get old regions names for unregister
 
         Args:
-            regions_statuses (list): responce of get_cuda_shared_memory_status from triton
+            regions_statuses (list): responce of get_shared_memory_status from triton
             new_triton_shm_name (str): name of new region
 
         Returns:
             List[str]: old regions names for unregister
         """
+        if self.max_shm_regions < 1:
+            return []
         i_sep = len(new_triton_shm_name) - new_triton_shm_name[::-1].index('_') - 1
         region_name = new_triton_shm_name[:i_sep]
         registrated_regions = [
@@ -231,44 +269,35 @@ class TritonClient(BaseClient):
             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
+    def _create_shm_handle(self, shape: List[int], dtype: np.dtype, name: str) -> ShmHandlerWrapper:
         """
-        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
+        Create shared memory handle
 
         Args:
-            shape (List[int]): Shape of cuda shared memory region
+            shape (List[int]): Shape of shared memory region
             dtype (np.dtype): Data type of input/output data
             name (str): Input/output name
 
         Returns:
-            Any: CUDA shared memory handle
+            Any: 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)
+        if self.use_cuda_shm:
+            shm_handle = cudashm.create_shared_memory_region(
+                shm_name, byte_size, 0
+                )
+        else:
+            shm_handle = shm.create_shared_memory_region(
+                shm_name, shm_name, byte_size
+                )
+        return ShmHandlerWrapper(shm_handle, shm_name, byte_size)
 
-    def _create_cuda_shm_handles_for_io(self, shapes: List[List[int]], 
+    def _create_shm_handles_for_io(self, shapes: List[List[int]], 
                                         dtypes: List[np.dtype], 
-                                        names: List[str]) -> List[Any]:
+                                        names: List[str]) -> List[ShmHandlerWrapper]:
         """
-        Create CUDA shared memory handles for inputs or outputs
+        Create shared memory handles for inputs or outputs
 
         Args:
             shapes (List[List[int]]): Shapes of cuda shared memory regions
@@ -276,24 +305,24 @@ class TritonClient(BaseClient):
             names (List[str]): Input/output names
 
         Returns:
-            List[Any]: CUDA shared memory handles
+            List[ShmHandlerWrapper]: shared memory handles
         """
-        return [self._create_cuda_shm_handle(shape, dtype, name) 
+        return [self._create_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
+        Create shared memory handles for inputs and outputs
         """
-        self.input_shm_handles = self._create_cuda_shm_handles_for_io(
+        self.input_shm_handlers = self._create_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.output_shm_handlers = self._create_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:
+                             config_input_format: str, shm_handler: Optional[ShmHandlerWrapper] = None) -> Any:
         """
         Create triton InferInput
 
@@ -301,27 +330,28 @@ class TritonClient(BaseClient):
             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.
+            shm_handler (ShmHandlerWrapper, optional): shared memory handler. 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)
+        if self.use_cuda_shm or self.use_system_shm:
+            shm_utils = cudashm if self.use_cuda_shm else shm
+            shm_utils.set_shared_memory_region(shm_handler.handler, [input_data])
+            infer_input.set_shared_memory(shm_handler.name, shm_handler.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:
+    def _create_triton_output(self, output_name: str, binary: bool = True, shm_handler: Optional[ShmHandlerWrapper] = 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.
+            shm_handler (ShmHandlerWrapper, optional): shared memory handler. Defaults to None.
 
         Returns:
             Any: triton InferRequestedOutput for receiving response
@@ -330,10 +360,39 @@ class TritonClient(BaseClient):
             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)
+        if self.use_cuda_shm or self.use_system_shm:
+            infer_output.set_shared_memory(shm_handler.name, shm_handler.size)
         return infer_output
 
+    def _register_shm_regions(self):
+        """
+        Register shared memory regions in Triton
+        """
+        get_shared_memory_status = self.triton_client.get_cuda_shared_memory_status \
+            if self.use_cuda_shm else self.triton_client.get_system_shared_memory_status
+        
+        unregister_shared_memory = self.triton_client.unregister_cuda_shared_memory \
+            if self.use_cuda_shm else self.triton_client.unregister_system_shared_memory
+        
+        if self.scheme == "grpc":
+            regions_statuses = get_shared_memory_status(as_json=True)['regions']
+        else:
+            regions_statuses = get_shared_memory_status()
+
+        for shm_handler in self.input_shm_handlers + self.output_shm_handlers:
+            old_regions_names = self._get_old_regions_names(regions_statuses, shm_handler.name)
+            for old_region_name in old_regions_names:
+                unregister_shared_memory(old_region_name)
+            
+            if self.use_cuda_shm:
+                self.triton_client.register_cuda_shared_memory(
+                    shm_handler.name, cudashm.get_raw_handle(shm_handler.handler), 0, shm_handler.size
+                )
+            else:
+                self.triton_client.register_system_shared_memory(
+                    shm_handler.name, shm_handler.name, shm_handler.size
+                )
+
     def _postprocess_triton_result(self, triton_response: Any, padding_size: int) -> Dict[str, np.ndarray]:
         """
         Postprocess triton response.
@@ -346,15 +405,17 @@ class TritonClient(BaseClient):
             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:
+        for output_name, shm_op_handle in zip(self.outputs_names, self.output_shm_handlers):
+            if self.use_cuda_shm or self.use_system_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,
+                    
+                shm_utils = shm if self.use_system_shm else cudashm
+                result[output_name] = shm_utils.get_contents_as_numpy(
+                    shm_op_handle.handler,
                     utils.triton_to_np_dtype(output["datatype"]),
                     output["shape"],
                 )
@@ -375,17 +436,17 @@ class TritonClient(BaseClient):
         
         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):
+            for input_name, config_input_format, shm_ip_handler in \
+                    zip(self.inputs_names, self.triton_inputs_dtypes, self.input_shm_handlers):
                 triton_input = self._create_triton_input(
-                    inputs_batches[input_name][i_batch], input_name, config_input_format, shm_ip_handle
+                    inputs_batches[input_name][i_batch], input_name, config_input_format, shm_ip_handler
                     )
                 triton_inputs.append(triton_input)
 
             triton_outputs = []
-            for output_name, shm_op_handle in zip(self.outputs_names, self.output_shm_handles):
+            for output_name, shm_op_handlers in zip(self.outputs_names, self.output_shm_handlers):
                 triton_output = self._create_triton_output(
-                    output_name, binary=True, shm_handle=shm_op_handle
+                    output_name, binary=True, shm_handler=shm_op_handlers
                     )
                 triton_outputs.append(triton_output)
 
@@ -413,14 +474,14 @@ class TritonClient(BaseClient):
         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):
+                    zip(self.inputs_names, self.triton_inputs_dtypes, self.input_shm_handlers):
                 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):
+            for output_name, shm_op_handle in zip(self.outputs_names, self.output_shm_handlers):
                 triton_output = self._create_triton_output(
                     output_name, binary=True, shm_handle=shm_op_handle
                     )

{imb-1.0.1.dist-info → imb-1.0.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: imb
-Version: 1.0.1
+Version: 1.0.2
 Summary: Python library for run inference of deep learning models in different backends
 Home-page: https://github.com/TheConstant3/InferenceMultiBackend
 Author: p-constant
@@ -55,6 +55,8 @@ For support all implemented clients:
 
 OnnxClient usage example
 ```
+from imb.onnx import OnnxClient
+
 onnx_client = OnnxClient(
     model_path='model.onnx',
     model_name='any name',
@@ -64,15 +66,19 @@ onnx_client = OnnxClient(
     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
 ```
+from imb.triton import TritonClient
+
 triton_client = TritonClient(
     url='localhost:8000',
     model_name='arcface',
@@ -87,9 +93,11 @@ triton_client = TritonClient(
     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()])
 ```

imb-1.0.2.dist-info/RECORD

@@ -0,0 +1,9 @@
+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=92d3tvCniWGSnC1UyjkQ5OcXgSbsBnX6T2hoewLal0k,21796
+imb-1.0.2.dist-info/LICENSE,sha256=pAZXnNE2dxxwXFIduGyn1gpvPefJtUYOYZOi3yeGG94,1068
+imb-1.0.2.dist-info/METADATA,sha256=lEzhVDdcdNHZeECQPisnQcZDjueOFP8zuhVTDh4Vi3s,3314
+imb-1.0.2.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+imb-1.0.2.dist-info/top_level.txt,sha256=kY8Fp1i_MzTZhuoVhVexG762D8HBd-THfX_lfw4EZmY,4
+imb-1.0.2.dist-info/RECORD,,

imb/inference_clients/base.py

@@ -1,129 +0,0 @@
-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/onnx.py

@@ -1,99 +0,0 @@
-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/inference_clients/triton.py

@@ -1,460 +0,0 @@
-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/RECORD

@@ -1,13 +0,0 @@
-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,,