inference_server.testing

Functions for testing inference-server plugins

class ImplementsDeserialize(*args, **kwargs)

Bases: Protocol

Interface compatible with sagemaker.deserializers.BaseDeserializer

property ACCEPT: Tuple[str]

The content types that are supported by this deserializer

deserialize(stream: StreamingBody, content_type: str) → Any

Return the deserialized data

class ImplementsSerialize(*args, **kwargs)

Bases: Protocol

Interface compatible with sagemaker.serializers.BaseSerializer

property CONTENT_TYPE: str

The MIME type for the serialized data

serialize(data: Any) → bytes

Return the serialized data

client() → Client

Return an HTTP test client for inference_server

The test client is simply a werkzeug.test.Client instance which loads the inference-server WSGI app. Consult the werkzeug documentation for details how to use the test client.

hookimpl_is_valid(function: Callable) → bool

Return whether the given function is a valid implementation of an inference_server hook

Parameters:

function – The hook function to validate

plugin_is_registered(plugin: Type | ModuleType) → bool

Return whether the given plugin is registered with inference_server

This validates whether a plugin entrypoint is defined in pyproject.toml like this:

[project.entry-points.inference_server]
my_plugin_name = "my_module_name"

Parameters:

plugin – The plugin, typically a module containg the hook functions.

plugin_manager() → PluginManager

Return the plugin manager used by inference-server

post_invocations(*, model_dir: Path | None = None, **kwargs) → TestResponse

Send an HTTP POST request to /invocations using a test HTTP client and return the response

This function should be used to verify an inference request using the full inference-server logic.

Parameters:

• model_dir – Optional pass a custom model directory to load the model from. Default is /opt/ml/model/.

• kwargs – Keyword arguments passed to werkzeug.test.Client.post()

predict(data: Any, *, model_dir: Path | None = None, serializer: ImplementsSerialize | None = None, deserializer: ImplementsDeserialize | None = None) → Any

Invoke the model and return a prediction

Parameters:

• data – Model input data

• model_dir – Optional pass a custom model directory to load the model from. Default is /opt/ml/model/.

• serializer – Optional. A serializer for sending the data as bytes to the model server. Should be compatible with sagemaker.serializers.BaseSerializer. Default: bytes pass-through.

• deserializer – Optional. A deserializer for processing the prediction as sent by the model server. Should be compatible with sagemaker.deserializers.BaseDeserializer. Default: bytes pass-through.