Add asyncio support for tritonclient (beta) (#23)

* Add aio_grpcclient, aio_httpclient
* add test_multiple_tasks test-case
* tritony==0.0.19

Author: kimdwkimdw

.github/workflows/pre-commit_pytest.yml

@@ -52,7 +52,8 @@ jobs:
           TRITON_CLIENT_TIMEOUT: 1
         run: |
           tritonserver --model-repo=$GITHUB_WORKSPACE/model_repository &
-          pip install .[tests]
+          pip install uv
+          uv pip install .[tests]
           sleep 3
 
           curl -v ${TRITON_HOST}:${TRITON_HTTP}/v2/health/ready

bin/run_triton_tritony_sample.sh

@@ -11,7 +11,7 @@ docker run -it --rm --name triton_tritony \
     -e OMP_NUM_THREADS=2 \
     -e OPENBLAS_NUM_THREADS=2 \
     --shm-size=1g  \
-    nvcr.io/nvidia/tritonserver:23.08-pyt-python-py3 \
+    nvcr.io/nvidia/tritonserver:24.05-pyt-python-py3 \
     tritonserver --model-repository=/models \
     --exit-timeout-secs 15 \
     --min-supported-compute-capability 7.0 \

model_repository/sample_sleep_1sec/1/model.py

@@ -0,0 +1,29 @@
+import json
+import time
+
+import triton_python_backend_utils as pb_utils
+
+
+class TritonPythonModel:
+    def initialize(self, args):
+        self.model_config = model_config = json.loads(args["model_config"])
+        output_configs = model_config["output"]
+
+        self.output_name_list = [output_config["name"] for output_config in output_configs]
+        self.output_dtype_list = [
+            pb_utils.triton_string_to_numpy(output_config["data_type"]) for output_config in output_configs
+        ]
+
+    def execute(self, requests):
+        responses = [None for _ in requests]
+        for idx, request in enumerate(requests):
+            current_add_value = int(json.loads(request.parameters()).get("add", 0))
+            in_tensor = [item.as_numpy() + current_add_value for item in request.inputs() if item.name() == "model_in"]
+            out_tensor = [
+                pb_utils.Tensor(output_name, x.astype(output_dtype))
+                for x, output_name, output_dtype in zip(in_tensor, self.output_name_list, self.output_dtype_list)
+            ]
+            inference_response = pb_utils.InferenceResponse(output_tensors=out_tensor)
+            responses[idx] = inference_response
+        time.sleep(1)
+        return responses

model_repository/sample_sleep_1sec/config.pbtxt

@@ -0,0 +1,40 @@
+backend: "python"
+max_batch_size: 0
+
+input [
+{
+    name: "model_in"
+    data_type: TYPE_FP32
+    dims: [ -1 ]
+}
+]
+
+output [
+{
+    name: "model_out"
+    data_type: TYPE_FP32
+    dims: [ -1 ]
+}
+]
+
+instance_group [{ kind: KIND_CPU, count: 10 }]
+
+model_warmup {
+  name: "RandomSampleInput"
+  batch_size: 1
+  inputs [{
+      key: "model_in"
+      value: {
+        data_type: TYPE_FP32
+        dims: [ 10 ]
+        random_data: true
+      }
+   }, {
+      key: "model_in"
+      value: {
+        data_type: TYPE_FP32
+        dims: [ 10 ]
+        zero_data: true
+      }
+   }]
+}

packaging.md

@@ -17,4 +17,7 @@ line-length = 120
 ignore = ["F811","F841","E203","E402","E501","E712","B019"]
 
 [tool.lint.isort]
-forced-separate = ["tests"]
+forced-separate = ["tests"]
+
+[tool.pytest.ini_options]
+asyncio_default_fixture_loop_scope = "function"

pyproject.toml

@@ -0,0 +1,3 @@
+[pytest]
+log_cli=true
+log_level=NOTSET

pytest.ini

@@ -69,6 +69,7 @@ tests =
     pytest-xdist
     pytest-mpl
     pytest-cov
+    pytest-asyncio
     pytest
     pre-commit
     coveralls

setup.cfg

@@ -1,7 +1,15 @@
+import logging
 import os
 
 import pytest
 
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
+    datefmt="%m/%d/%Y %H:%M:%S",
+)
+
+
 MODEL_NAME = os.environ.get("MODEL_NAME", "sample")
 TRITON_HOST = os.environ.get("TRITON_HOST", "localhost")
 TRITON_HTTP = os.environ.get("TRITON_HTTP", "8100")
@@ -14,3 +22,11 @@ def config(request):
     Returns a tuple of (protocol, port, run_async)
     """
     return request.param
+
+
+@pytest.fixture(params=[("http", TRITON_HTTP), ("grpc", TRITON_GRPC)])
+def async_config(request):
+    """
+    Returns a tuple of (protocol, port)
+    """
+    return request.param

tests/common_fixtures.py

@@ -0,0 +1,58 @@
+import asyncio
+import logging
+
+import numpy as np
+import pytest
+
+from tritony import InferenceClient
+
+from .common_fixtures import MODEL_NAME, TRITON_HOST, async_config
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
+    datefmt="%m/%d/%Y %H:%M:%S",
+)
+
+__all__ = ["async_config"]
+EPSILON = 1e-8
+
+
+def get_client(protocol, port, model_name):
+    print(f"Testing {protocol}", flush=True)
+    return InferenceClient.create_with_asyncio(model_name, f"{TRITON_HOST}:{port}", protocol=protocol)
+
+
+@pytest.mark.asyncio
+async def test_basics(async_config):
+    protocol, port = async_config
+
+    client = get_client(*async_config, model_name=MODEL_NAME)
+    sample = np.random.rand(1, 100).astype(np.float32)
+
+    result = await client.aio_infer(sample)
+    assert np.isclose(result, sample).all()
+
+    result = await client.aio_infer({"model_in": sample})
+    assert np.isclose(result, sample).all()
+
+
+@pytest.mark.asyncio
+async def test_multiple_tasks(async_config):
+    n_multiple_tasks = 10
+    protocol, port = async_config
+    print(f"Testing {protocol}:{port}")
+
+    client_list = [get_client(*async_config, model_name="sample_sleep_1sec") for _ in range(n_multiple_tasks)]
+
+    sample = np.random.rand(1, 100).astype(np.float32)
+    tasks = [client.aio_infer(sample) for client in client_list]
+
+    start_time = asyncio.get_event_loop().time()
+    results = await asyncio.gather(*tasks)
+    end_time = asyncio.get_event_loop().time()
+
+    for result in results:
+        assert np.isclose(result, sample).all()
+
+    assert (end_time - start_time) < 2, f"Time taken: {end_time - start_time}"