Add TritonModelInput with optional

Author: kimdwkimdw

.pre-commit-config.yaml

@@ -28,7 +28,7 @@ repos:
     hooks:
       - id: flake8
         types: [python]
-        args: ["--max-line-length", "120", "--ignore", "F811,F841,E203,E402,E712,W503"]
+        args: ["--max-line-length", "120", "--ignore", "F811,F841,E203,E402,E712,W503,E501"]
   - repo: https://github.com/shellcheck-py/shellcheck-py
     rev: v0.9.0.5
     hooks:

setup.cfg

@@ -43,6 +43,8 @@ classifiers =
 zip_safe = False
 include_package_data = True
 packages = find:
+package_dir =
+    =.
 install_requires =
     tritonclient[all]>=2.21.0
     protobuf>=3.5.0

tritony/helpers.py

@@ -10,6 +10,7 @@
 from attrs import define
 from tritonclient import grpc as grpcclient
 from tritonclient import http as httpclient
+from tritonclient.grpc import model_config_pb2
 
 
 class TritonProtocol(Enum):
@@ -31,13 +32,32 @@ def dict_to_attr(obj: dict[str, Any]) -> SimpleNamespace:
     return json.loads(json.dumps(obj), object_hook=lambda d: SimpleNamespace(**d))
 
 
+@define
+class TritonModelInput:
+    """
+    Most of the fields are mapped to model_config_pb2.ModelInput(https://github.com/triton-inference-server/common/blob/a2de06f4c80b2c7b15469fa4d36e5f6445382bad/protobuf/model_config.proto#L317)
+
+    Commented fields are not used.
+    """
+
+    name: str
+    dtype: str  # data_type mapping to https://github.com/triton-inference-server/client/blob/d257c0e5c3de6e15d6ef289ff2b96cecd0a69b5f/src/python/library/tritonclient/utils/__init__.py#L163-L190
+
+    format: int = 0
+    dims: list[int] = []  # dims
+
+    # reshape: list[int] = []
+    # is_shape_tensor: bool = False
+    # allow_ragged_batch: bool = False
+    optional: bool = False
+
+
 @define
 class TritonModelSpec:
     name: str
 
     max_batch_size: int
-    input_name: list[str]
-    input_dtype: list[str]
+    model_input: list[TritonModelInput]
 
     output_name: list[str]
 
@@ -91,7 +111,7 @@ def get_triton_client(
     model_name: str,
     model_version: str,
     protocol: TritonProtocol,
-):
+) -> (int, list[TritonModelInput], list[str]):
     """
     (required in)
     :param triton_client:
@@ -107,23 +127,43 @@ def get_triton_client(
 
     args = dict(model_name=model_name, model_version=model_version)
 
-    model_metadata = triton_client.get_model_metadata(**args)
     model_config = triton_client.get_model_config(**args)
     if protocol is TritonProtocol.http:
-        model_metadata = dict_to_attr(model_metadata)
         model_config = dict_to_attr(model_config)
     elif protocol is TritonProtocol.grpc:
         model_config = model_config.config
 
-    max_batch_size, input_name_list, output_name_list, dtype_list = parse_model(model_metadata, model_config)
+    max_batch_size, input_list, output_name_list = parse_model(model_config)
+
+    return max_batch_size, input_list, output_name_list
+
 
-    return max_batch_size, input_name_list, output_name_list, dtype_list
+def parse_model_input(
+    model_input: model_config_pb2.ModelInput | SimpleNamespace,
+) -> TritonModelInput:
+    """
+    https://github.com/triton-inference-server/common/blob/r23.08/protobuf/model_config.proto#L317-L412
+    """
+    RAW_DTYPE = model_input.data_type
+    if isinstance(model_input.data_type, int):
+        RAW_DTYPE = model_config_pb2.DataType.Name(RAW_DTYPE)
+    RAW_DTYPE = RAW_DTYPE.strip("TYPE_")
+
+    if RAW_DTYPE == "STRING":
+        RAW_DTYPE = "BYTES"  # https://github.com/triton-inference-server/client/blob/d257c0e5c3de6e15d6ef289ff2b96cecd0a69b5f/src/python/library/tritonclient/utils/__init__.py#L188-L189
+    return TritonModelInput(
+        name=model_input.name,
+        dims=model_input.dims,
+        dtype=RAW_DTYPE,
+        optional=model_input.optional,
+    )
 
 
-def parse_model(model_metadata, model_config):
+def parse_model(
+    model_config: model_config_pb2.ModelConfig | SimpleNamespace,
+) -> (int, list[TritonModelInput], list[str]):
     return (
         model_config.max_batch_size,
-        [input_metadata.name for input_metadata in model_metadata.inputs],
-        [output_metadata.name for output_metadata in model_metadata.outputs],
-        [input_metadata.datatype for input_metadata in model_metadata.inputs],
+        [parse_model_input(model_config_input) for model_config_input in model_config.input],
+        [model_config_output.name for model_config_output in model_config.output],
     )

tritony/tools.py

@@ -6,7 +6,6 @@
 import logging
 import os
 import time
-import warnings
 from concurrent.futures import ThreadPoolExecutor
 from typing import Any, Dict, List, Optional, Union
 
@@ -198,14 +197,6 @@ def triton_client(self):
     def default_model_spec(self):
         return self.model_specs[self.default_model]
 
-    @property
-    def input_name_list(self):
-        warnings.warn(
-            "input_name_list is deprecated, please use 'default_model_spec.input_name' instead", DeprecationWarning
-        )
-
-        return self.default_model_spec.input_name
-
     def __del__(self):
         # Not supporting streaming
         # if self.flag.protocol is TritonProtocol.grpc and self.flag.streaming and hasattr(self, "triton_client"):
@@ -223,15 +214,14 @@ def _renew_triton_client(self, triton_client, model_name: str | None = None, mod
         triton_client.is_server_ready()
         triton_client.is_model_ready(model_name, model_version)
 
-        (max_batch_size, input_name_list, output_name_list, dtype_list) = get_triton_client(
+        (max_batch_size, input_list, output_name_list) = get_triton_client(
             triton_client, model_name=model_name, model_version=model_version, protocol=self.flag.protocol
         )
 
         self.model_specs[(model_name, model_version)] = TritonModelSpec(
             name=model_name,
             max_batch_size=max_batch_size,
-            input_name=input_name_list,
-            input_dtype=dtype_list,
+            model_input=input_list,
             output_name=output_name_list,
         )
 
@@ -257,7 +247,12 @@ def __call__(
         if type(sequences_or_dict) in [list, np.ndarray]:
             sequences_list = [sequences_or_dict]
         elif type(sequences_or_dict) is dict:
-            sequences_list = [sequences_or_dict[input_name] for input_name in model_spec.input_name]
+            sequences_list = [
+                sequences_or_dict[model_input.name]
+                for model_input in model_spec.model_input
+                if model_input.optional is False  # check required
+                or (model_input.optional is True and model_input.name in sequences_or_dict)  # check optional
+            ]
 
         return self._call_async(sequences_list, model_spec=model_spec)
 
@@ -267,8 +262,8 @@ def build_triton_input(self, _input_list: List[np.array], model_spec: TritonMode
         else:
             client = httpclient
         infer_input_list = []
-        for _input, _input_name, _dtype in zip(_input_list, model_spec.input_name, model_spec.input_dtype):
-            infer_input = client.InferInput(_input_name, _input.shape, _dtype)
+        for _input, _model_input in zip(_input_list, model_spec.model_input):
+            infer_input = client.InferInput(_model_input.name, _input.shape, _model_input.dtype)
             infer_input.set_data_from_numpy(_input)
             infer_input_list.append(infer_input)