aiSSEMBLE™ Open Inference Protocol gRPC

The Open Inference Protocol (OIP) specification defines a standard protocol for performing machine learning model inference across serving runtimes for different ML frameworks. This Python application can be leveraged to create a gRPC server that is compatible with the Open Inference Protocol. It handles standing up and tearing down the server so you only need to worry about the inferencing functionality.

Installation

Add aissemble-open-inference-protocol-grpc to an application

pip install aissemble-open-inference-protocol-grpc

Usage

Creating the Server

Use aissemble-open-inference-protocol-grpc to create a gRPC server by creating a file main.py with:

import asyncio
from aissemble_open_inference_protocol_grpc.aissemble_oip_grpc import AissembleOIPgRPC

grpc = AissembleOIPgRPC()

if __name__ == '__main__':
    asyncio.run(grpc.start_server())

The gRPC server will come up after a few seconds and will be OIP compliant. The proto specifications can be found in the grpc_inference_service.proto file.

Implementing the Endpoints Handler

By default, most of the gRPC endpoints will return a Method Not Implemented. You can implement these functions by creating a custom handler extending DataplaneHandler.

Example:

from typing import Optional

from aissemble_open_inference_protocol_shared.handlers.model_handler import (
    ModelHandler,
)
from aissemble_open_inference_protocol_shared.types.dataplane import (
    Datatype,
    InferenceRequest,
    InferenceResponse,
    ModelMetadataResponse,
    MetadataTensor,
)


class MyHandler(ModelHandler):
    def __init__(self):
        super().__init__()

    def infer(
            self,
            payload: InferenceRequest,
            model_name: str,
            model_version: Optional[str] = None,
    ) -> InferenceResponse:
        return InferenceResponse(
            model_name=model_name, model_version=model_version, id="id", outputs=[]
        )

    def model_metadata(
            self,
            model_name: str,
            model_version: Optional[str] = None,
    ) -> ModelMetadataResponse:
        return ModelMetadataResponse(
            name=model_name,
            versions=[model_version] if model_version else None,
            platform="python",
            inputs=[MetadataTensor(name="input", datatype=Datatype.FP32, shape=[1])],
            outputs=[
                MetadataTensor(name="output", datatype=Datatype.FP32, shape=[1])
            ],
        )

    def model_load(self, model_name: str) -> bool:
        # Do some model loading
        return True

Use aissemble-open-inference-protocol-grpc to create a gRPC server and pass it MyHandler

from aissemble_open_inference_protocol_grpc.aissemble_oip_grpc import AissembleOIPgRPC

grpc = AissembleOIPgRPC(MyHandler())

Now when starting the server, the inference requests will route to the handler.

Configuration

There are several configurations available that affect the server. These can be implemented via Krausening properties file oip.properties or environment variables.

Examples

For working examples, refer to the Examples documentation.