How to enable server-side batching on SageMaker PyTorch TorchServe endpoints?

In TorchServe itself, there are two configuration parameter that setup the server-side batching

feature:

• batchSize: This is the maximum batch size in ms that a model is expected to handle.
• maxBatchDelay: This is the maximum batch delay time TorchServe waits to receive batch_size number of requests. If TorchServe doesn’t receive batch_size number of requests before this timer time’s out, it sends what ever requests that were received to the model handler.

With the recent update of PyTorch Inference Toolkit for SageMaker (October 2021), these variables have been exposed to user via environmental variables:

• batchSize is set by SAGEMAKER_TS_BATCH_SIZE
• maxBatchDelay is set by SAGEMAKER_TS_MAX_BATCH_DELAY

All in all, an example how to set up server-side batching using SageMaker Python SDK:

from sagemaker.pytorch.model import PyTorchModel

env_variables_dict = {"SAGEMAKER_TS_BATCH_SIZE": "3","SAGEMAKER_TS_MAX_BATCH_DELAY": "100000"}

pytorch_model = PyTorchModel(
    model_data=model_artifact,
    role=role,
    image_uri=image_uri,
    source_dir="code",
    framework_version='1.9',
    entry_point="inference.py",
    env=env_variables_dict)

See also this AWS ML blog post:

Optimize your inference jobs using dynamic batch inference with TorchServe on Amazon SageMaker