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

imb/__init__.py

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

imb/{inference_clients/triton.py → 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.2.dist-info/METADATA

@@ -0,0 +1,113 @@
+Metadata-Version: 2.2
+Name: imb
+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
+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
+```
+from imb.onnx import OnnxClient
+
+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
+```
+from imb.triton import TritonClient
+
+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.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/__init__.py

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

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,,