Which SageMaker server supports server-side batching and how to enable it?

Francesco Pochetti 0
Viewed 17

MMS, TFServing and TorchServe support server-side batching (consequent requests can be locally batched in async fashion by the server while maintaining the illusion of synchronous batch-1 size to the client). How to enable those features on SageMaker endpoints?

amazon-sagemaker
1 Answers

For each SageMaker container these can be controlled via environment variables.

For TorchServe:

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
)

ENVs in toolkit: https://github.com/aws/sagemaker-pytorch-inference-toolkit/blob/27b667fa27259dcea92b97e3dcc903057587deb6/src/sagemaker_pytorch_serving_container/ts_parameters.py

Blog post for more information: https://aws.amazon.com/blogs/machine-learning/optimize-your-inference-jobs-using-dynamic-batch-inference-with-torchserve-on-amazon-sagemaker/

TFServing Batching docs: https://github.com/aws/sagemaker-tensorflow-serving-container/blob/1bd309b7be5040d5515a3081fd5714e444b2ab91/README.md#enabling-batching

MMS batching feature is not supported in SageMaker right now.

Marc K0
Related