|
|
|
"""Utils to interact with the Triton Inference Server.""" |
|
|
|
import typing |
|
from urllib.parse import urlparse |
|
|
|
import torch |
|
|
|
|
|
class TritonRemoteModel: |
|
""" |
|
A wrapper over a model served by the Triton Inference Server. |
|
|
|
It can be configured to communicate over GRPC or HTTP. It accepts Torch Tensors as input and returns them as |
|
outputs. |
|
""" |
|
|
|
def __init__(self, url: str): |
|
""" |
|
Keyword arguments: |
|
url: Fully qualified address of the Triton server - for e.g. grpc://localhost:8000 |
|
""" |
|
|
|
parsed_url = urlparse(url) |
|
if parsed_url.scheme == "grpc": |
|
from tritonclient.grpc import InferenceServerClient, InferInput |
|
|
|
self.client = InferenceServerClient(parsed_url.netloc) |
|
model_repository = self.client.get_model_repository_index() |
|
self.model_name = model_repository.models[0].name |
|
self.metadata = self.client.get_model_metadata(self.model_name, as_json=True) |
|
|
|
def create_input_placeholders() -> typing.List[InferInput]: |
|
return [ |
|
InferInput(i["name"], [int(s) for s in i["shape"]], i["datatype"]) for i in self.metadata["inputs"] |
|
] |
|
|
|
else: |
|
from tritonclient.http import InferenceServerClient, InferInput |
|
|
|
self.client = InferenceServerClient(parsed_url.netloc) |
|
model_repository = self.client.get_model_repository_index() |
|
self.model_name = model_repository[0]["name"] |
|
self.metadata = self.client.get_model_metadata(self.model_name) |
|
|
|
def create_input_placeholders() -> typing.List[InferInput]: |
|
return [ |
|
InferInput(i["name"], [int(s) for s in i["shape"]], i["datatype"]) for i in self.metadata["inputs"] |
|
] |
|
|
|
self._create_input_placeholders_fn = create_input_placeholders |
|
|
|
@property |
|
def runtime(self): |
|
"""Returns the model runtime.""" |
|
return self.metadata.get("backend", self.metadata.get("platform")) |
|
|
|
def __call__(self, *args, **kwargs) -> typing.Union[torch.Tensor, typing.Tuple[torch.Tensor, ...]]: |
|
""" |
|
Invokes the model. |
|
|
|
Parameters can be provided via args or kwargs. args, if provided, are assumed to match the order of inputs of |
|
the model. kwargs are matched with the model input names. |
|
""" |
|
inputs = self._create_inputs(*args, **kwargs) |
|
response = self.client.infer(model_name=self.model_name, inputs=inputs) |
|
result = [] |
|
for output in self.metadata["outputs"]: |
|
tensor = torch.as_tensor(response.as_numpy(output["name"])) |
|
result.append(tensor) |
|
return result[0] if len(result) == 1 else result |
|
|
|
def _create_inputs(self, *args, **kwargs): |
|
args_len, kwargs_len = len(args), len(kwargs) |
|
if not args_len and not kwargs_len: |
|
raise RuntimeError("No inputs provided.") |
|
if args_len and kwargs_len: |
|
raise RuntimeError("Cannot specify args and kwargs at the same time") |
|
|
|
placeholders = self._create_input_placeholders_fn() |
|
if args_len: |
|
if args_len != len(placeholders): |
|
raise RuntimeError(f"Expected {len(placeholders)} inputs, got {args_len}.") |
|
for input, value in zip(placeholders, args): |
|
input.set_data_from_numpy(value.cpu().numpy()) |
|
else: |
|
for input in placeholders: |
|
value = kwargs[input.name] |
|
input.set_data_from_numpy(value.cpu().numpy()) |
|
return placeholders |
|
|
