Principle:Triton inference server Server SageMaker Integration

Overview

SageMaker Integration is the principle governing how Triton Inference Server presents an AWS SageMaker-compatible HTTP endpoint, enabling Triton to be deployed as a SageMaker inference container without modification to client code or SageMaker infrastructure. The SagemakerAPIServer class extends Triton's base HTTP API server to implement SageMaker's proprietary endpoint contract -- /ping for health checks, /invocations for single-model inference, and /models for Multi-Model Endpoint (MME) lifecycle management -- while internally delegating to Triton's standard inference pipeline.

Theoretical Basis

Why SageMaker Compatibility Matters

AWS SageMaker is one of the most widely used managed ML platforms. It prescribes a specific HTTP contract that inference containers must implement: a health check endpoint on /ping, an inference endpoint on /invocations, and (for MME) model management endpoints under /models. By implementing this contract natively, Triton can serve as a drop-in SageMaker container, allowing users to leverage Triton's multi-framework support, dynamic batching, model ensembles, and GPU optimization while remaining within the SageMaker managed infrastructure.

Endpoint Mapping

SageMaker Route	Triton Internal Operation	HTTP Method
/ping	Server health (live/ready)	GET
/invocations	Inference on default model	POST
/models/{name}/invoke	Inference on named model (MME)	POST
/models/{name} (POST)	Load model from repository (MME)	POST
/models/{name} (DELETE)	Unload model (MME)	DELETE
/models/{name} (GET)	Get model status (MME)	GET
/models (GET)	List loaded models (MME)	GET

Environment Variable Configuration

SageMaker containers are configured primarily through environment variables rather than CLI flags. The server reads:

• SAGEMAKER_TRITON_DEFAULT_MODEL_NAME: The model name for /invocations requests in single-model mode.
• SAGEMAKER_TRITON_PING_MODE: Controls whether /ping reports "live" or "ready" status (default: "ready").
• SAGEMAKER_TRITON_INFERENCE_TYPE: Determines whether /invocations maps to "infer", "generate", or "generate_stream", enabling LLM inference via the SageMaker endpoint.

Multi-Model Endpoint (MME) Support

SageMaker MME allows a single endpoint to host many models, loading and unloading them on demand. The SagemakerAPIServer maintains an internal sagemaker_models_list_ map (protected by a mutex for concurrent access) that tracks loaded models. When a load request arrives:

1. The server parses the url field from the JSON body to extract the model repository path (expected format: /opt/ml/models/{name}/model/{version}).
2. It adds the path as a model repository via TRITONSERVER_ServerRegisterModelRepository.
3. It loads the model via TRITONSERVER_ServerLoadModel.
4. On success, it registers the model in the local tracking map.

OOM Error Handling

SageMaker has specific expectations for how containers report out-of-memory conditions during model loading. The SageMakerMMECheckOOMError method inspects error messages for GPU memory exhaustion indicators, and SageMakerMMEHandleOOMError returns an HTTP 507 (Insufficient Storage) response with an appropriate error message, matching SageMaker's expected error contract.

Inference Header Differences

SageMaker's inference protocol does not use the KFServing Inference-Header-Content-Length header. The server overrides GetInferenceHeaderLength() to treat the entire request body as the inference payload, and overrides SetResponseHeader() to omit the binary data header that the standard KFServing protocol includes.

Compression Pass-Through

The SageMaker endpoint overrides both GetRequestCompressionType() and GetResponseCompressionType() to return IDENTITY, as the SageMaker compression protocol is not yet defined. This ensures no unexpected compression behavior when deployed within SageMaker infrastructure.

Unload with Timeout

Model unloading is asynchronous in Triton, so the SageMaker server polls model availability status in a loop with a 500ms sleep interval and a 350-second timeout, waiting until the model reports as UNAVAILABLE with reason "unloaded" before returning success to SageMaker.

Related Pages

Implementation:Triton_inference_server_Server_SagemakerServer Triton_inference_server_Server