"""AI wrapper over olo.ai.v1.Ai — generic ONNX model inference."""
import grpc
import numpy as np
from olo_protos.ai.v1 import ai_pb2
from olo.ai._tensors import numpy_dtype_for, numpy_to_tensor, tensor_to_numpy
from olo.ai.types import Detection, ModelInfo, TensorInfo
from olo.errors import wrap_rpc_error
_DEFAULT_TIMEOUT_SEC = 30.0
_DEFAULT_INFER_TIMEOUT_SEC = 60.0
def _tensor_info_from_proto(info: ai_pb2.TensorInfo) -> TensorInfo:
return TensorInfo(
name=info.name,
dtype=numpy_dtype_for(info.dtype),
shape=tuple(int(dim) for dim in info.shape),
)
def _model_info_from_proto(info: ai_pb2.ModelInfo) -> ModelInfo:
return ModelInfo(
model_id=info.model_id,
name=info.name,
inputs=tuple(_tensor_info_from_proto(item) for item in info.inputs),
outputs=tuple(_tensor_info_from_proto(item) for item in info.outputs),
metadata=dict(info.metadata),
)
[docs]
class Model:
"""Handle for a loaded ONNX model session on the appliance."""
[docs]
def __init__(self, client: "Ai", info: ModelInfo) -> None:
"""Bind to the Ai client and a loaded model's metadata."""
self._client = client
self._info = info
@property
def id(self) -> str:
"""Server-assigned session id."""
return self._info.model_id
@property
def info(self) -> ModelInfo:
"""Model input/output metadata."""
return self._info
[docs]
def infer(
self,
inputs: dict[str, np.ndarray],
*,
output_names: list[str] | None = None,
timeout: float = _DEFAULT_INFER_TIMEOUT_SEC,
) -> dict[str, np.ndarray]:
"""Run inference with named input arrays; returns named output arrays."""
return self._client._infer(self.id, inputs, output_names=output_names, timeout=timeout)
[docs]
def unload(self, *, timeout: float = _DEFAULT_TIMEOUT_SEC) -> None:
"""Release the model session on the appliance."""
self._client._unload(self.id, timeout=timeout)
def __enter__(self) -> "Model":
"""Enter a context manager."""
return self
def __exit__(self, exc_type, exc, tb) -> None:
"""Unload the model on context exit."""
self.unload()
[docs]
class Ai:
"""Sync wrapper for the AI gRPC service."""
[docs]
def __init__(self, session) -> None:
"""Bind to a channel session."""
self._session = session
@property
def _stub(self):
return self._session.ai_stub
[docs]
def list_models(self, *, timeout: float = _DEFAULT_TIMEOUT_SEC) -> list[str]:
"""List loadable model names in the appliance model directory."""
try:
response = self._stub.ListModels(ai_pb2.ListModelsRequest(), timeout=timeout)
except grpc.RpcError as exc:
raise wrap_rpc_error(exc) from exc
return list(response.names)
[docs]
def load_model(self, name: str, *, timeout: float = _DEFAULT_TIMEOUT_SEC) -> Model:
"""Load a model by file name and return a session handle."""
try:
response = self._stub.LoadModel(ai_pb2.LoadModelRequest(name=name), timeout=timeout)
except grpc.RpcError as exc:
raise wrap_rpc_error(exc) from exc
return Model(self, _model_info_from_proto(response.info))
[docs]
def model_info(self, model_id: str, *, timeout: float = _DEFAULT_TIMEOUT_SEC) -> ModelInfo:
"""Fetch metadata for an already-loaded model session."""
try:
response = self._stub.GetModelInfo(
ai_pb2.GetModelInfoRequest(model_id=model_id), timeout=timeout
)
except grpc.RpcError as exc:
raise wrap_rpc_error(exc) from exc
return _model_info_from_proto(response.info)
def _infer(
self,
model_id: str,
inputs: dict[str, np.ndarray],
*,
output_names: list[str] | None,
timeout: float,
) -> dict[str, np.ndarray]:
request = ai_pb2.InferRequest(model_id=model_id, output_names=output_names or [])
for name, array in inputs.items():
request.inputs[name].CopyFrom(numpy_to_tensor(np.asarray(array)))
try:
response = self._stub.Infer(request, timeout=timeout)
except grpc.RpcError as exc:
raise wrap_rpc_error(exc) from exc
return {name: tensor_to_numpy(tensor) for name, tensor in response.outputs.items()}
def _unload(self, model_id: str, *, timeout: float) -> None:
try:
self._stub.UnloadModel(ai_pb2.UnloadModelRequest(model_id=model_id), timeout=timeout)
except grpc.RpcError as exc:
raise wrap_rpc_error(exc) from exc
__all__ = [
"Ai",
"Detection",
"Model",
"ModelInfo",
"TensorInfo",
]