Unify definition of DeviceOrMemoryType

Author: kthui

python/tritonserver/_api/_datautils.py

@@ -35,6 +35,7 @@
 from typing import Any, Callable, ClassVar, Optional, Sequence, Type
 
 import numpy
+from tritonserver._api._memorybuffer import DeviceOrMemoryType
 from tritonserver._c import InvalidArgumentError
 from tritonserver._c import TRITONSERVER_DataType as DataType
 from tritonserver._c import TRITONSERVER_MemoryType as MemoryType
@@ -47,10 +48,6 @@
 except ImportError:
     cupy = None
 
-DeviceOrMemoryType = (
-    tuple[MemoryType, int] | MemoryType | tuple[_dlpack.DLDeviceType, int] | str
-)
-
 
 class CustomKeyErrorDict(dict):
     def __init__(

python/tritonserver/_api/_memorybuffer.py

@@ -33,12 +33,17 @@
 
 import numpy
 from tritonserver._api._datautils import DLPackObject
+from tritonserver._api._dlpack import DLDeviceType
 from tritonserver._c.triton_bindings import (
     InvalidArgumentError,
     TRITONSERVER_BufferAttributes,
 )
 from tritonserver._c.triton_bindings import TRITONSERVER_MemoryType as MemoryType
 
+DeviceOrMemoryType = (
+    tuple[MemoryType, int] | MemoryType | tuple[DLDeviceType, int] | str
+)
+
 
 @dataclass
 class MemoryBuffer:

python/tritonserver/_api/_request.py

@@ -36,17 +36,14 @@
 from tritonserver._api import _model
 from tritonserver._api._datautils import CustomKeyErrorDict
 from tritonserver._api._dlpack import DLDeviceType as DLDeviceType
+from tritonserver._api._memorybuffer import DeviceOrMemoryType
 from tritonserver._api._tensor import Tensor
 from tritonserver._c.triton_bindings import InvalidArgumentError
 from tritonserver._c.triton_bindings import TRITONSERVER_DataType as DataType
 from tritonserver._c.triton_bindings import TRITONSERVER_InferenceRequest
 from tritonserver._c.triton_bindings import TRITONSERVER_MemoryType as MemoryType
 from tritonserver._c.triton_bindings import TRITONSERVER_Server
 
-DeviceOrMemoryType = (
-    tuple[MemoryType, int] | MemoryType | tuple[DLDeviceType, int] | str
-)
-
 
 @dataclass
 class InferenceRequest:

python/tritonserver/_api/_response.py

@@ -43,7 +43,7 @@
 
 from tritonserver._api._dlpack import DLDeviceType as DLDeviceType
 from tritonserver._api._logging import LogMessage
-from tritonserver._api._memorybuffer import MemoryBuffer
+from tritonserver._api._memorybuffer import DeviceOrMemoryType, MemoryBuffer
 from tritonserver._api._tensor import Tensor
 from tritonserver._c.triton_bindings import (
     InternalError,
@@ -57,10 +57,6 @@
     TRITONSERVER_Server,
 )
 
-DeviceOrMemoryType = (
-    tuple[MemoryType, int] | MemoryType | tuple[DLDeviceType, int] | str
-)
-
 
 @dataclass
 class InferenceResponse:

python/tritonserver/_api/_tensor.py

@@ -45,7 +45,7 @@
     DLManagedTensor,
     c_str_dltensor,
 )
-from tritonserver._api._memorybuffer import MemoryBuffer
+from tritonserver._api._memorybuffer import DeviceOrMemoryType, MemoryBuffer
 from tritonserver._c.triton_bindings import (
     InvalidArgumentError,
     TRITONSERVER_BufferAttributes,
@@ -54,10 +54,6 @@
 from tritonserver._c.triton_bindings import TRITONSERVER_MemoryType as MemoryType
 from tritonserver._c.triton_bindings import UnsupportedError
 
-DeviceOrMemoryType = (
-    tuple[MemoryType, int] | MemoryType | tuple[DLDeviceType, int] | str
-)
-
 try:
     import cupy
 except ImportError: