Skip to content

Amazon SageMaker

Amazon SageMaker is a fully managed machine learning service that provides infrastructure and tools for building, training, and deploying ML models at scale. The Strands Agents SDK implements a SageMaker provider, allowing you to run agents against models deployed on SageMaker inference endpoints, including both pre-trained models from SageMaker JumpStart and custom fine-tuned models. The provider is designed to work with models that support OpenAI-compatible chat completion APIs.

For example, you can expose models like Mistral-Small-24B-Instruct-2501 on SageMaker, which has demonstrated reliable performance for conversational AI and tool calling scenarios.

SageMaker is configured as an optional dependency in Strands Agents. To install, run:

Terminal window
pip install 'strands-agents[sagemaker]' strands-agents-tools

After installing the SageMaker dependencies, you can import and initialize the Strands Agents’ SageMaker provider as follows:

from strands import Agent
from strands.models.sagemaker import SageMakerAIModel
from strands_tools import calculator
model = SageMakerAIModel(
endpoint_config={
"endpoint_name": "my-llm-endpoint",
"region_name": "us-west-2",
},
payload_config={
"max_tokens": 1000,
"temperature": 0.7,
"stream": True,
}
)
agent = Agent(model=model, tools=[calculator])
response = agent("What is the square root of 64?")

Note: Tool calling support varies by model. Models like Mistral-Small-24B-Instruct-2501 have demonstrated reliable tool calling capabilities, but not all models deployed on SageMaker support this feature. Verify your model’s capabilities before implementing tool-based workflows.

The endpoint_config configures the SageMaker endpoint connection:

ParameterDescriptionRequiredExample
endpoint_nameName of the SageMaker endpointYes"my-llm-endpoint"
region_nameAWS region where the endpoint is deployedYes"us-west-2"
inference_component_nameName of the inference componentNo"my-component"
target_modelSpecific model to invoke (multi-model endpoints)No"model-a.tar.gz"
target_variantProduction variant to invokeNo"variant-1"

The payload_config configures the model inference parameters:

ParameterDescriptionDefaultExample
max_tokensMaximum number of tokens to generateRequired1000
streamEnable streaming responsesTrueTrue
temperatureSampling temperature (0.0 to 2.0)Optional0.7
top_pNucleus sampling parameter (0.0 to 1.0)Optional0.9
top_kTop-k sampling parameterOptional50
stopList of stop sequencesOptional["Human:", "AI:"]

The SageMaker provider is designed to work with models that support OpenAI-compatible chat completion APIs. During development and testing, the provider has been validated with Mistral-Small-24B-Instruct-2501, which demonstrated reliable performance across various conversational AI tasks.

  • Model Performance: Results and capabilities vary significantly depending on the specific model deployed to your SageMaker endpoint
  • Tool Calling Support: Not all models deployed on SageMaker support function/tool calling. Verify your model’s capabilities before implementing tool-based workflows
  • API Compatibility: Ensure your deployed model accepts and returns data in the OpenAI chat completion format

For optimal results, we recommend testing your specific model deployment with your use case requirements before production deployment.

If you encounter ModuleNotFoundError: No module named 'boto3' or similar, install the SageMaker dependencies:

Terminal window
pip install 'strands-agents[sagemaker]'

The SageMaker provider uses standard AWS authentication methods (credentials file, environment variables, IAM roles, or AWS SSO). Ensure your AWS credentials have the necessary SageMaker invoke permissions.

Ensure your deployed model supports OpenAI-compatible chat completion APIs and verify tool calling capabilities if needed. Refer to the Model Compatibility section above for detailed requirements and testing recommendations.