AWS Sagemaker

LiteLLM supports All Sagemaker Huggingface Jumpstart Models

tip

We support ALL Sagemaker models, just set model=sagemaker/<any-model-on-sagemaker> as a prefix when sending litellm requests

API KEYS

os.environ["AWS_ACCESS_KEY_ID"] = ""
os.environ["AWS_SECRET_ACCESS_KEY"] = ""
os.environ["AWS_REGION_NAME"] = ""

Usage

import os 
from litellm import completion

os.environ["AWS_ACCESS_KEY_ID"] = ""
os.environ["AWS_SECRET_ACCESS_KEY"] = ""
os.environ["AWS_REGION_NAME"] = ""

response = completion(
            model="sagemaker/<your-endpoint-name>", 
            messages=[{ "content": "Hello, how are you?","role": "user"}],
            temperature=0.2,
            max_tokens=80
        )

Usage - Streaming

Sagemaker currently does not support streaming - LiteLLM fakes streaming by returning chunks of the response string

import os 
from litellm import completion

os.environ["AWS_ACCESS_KEY_ID"] = ""
os.environ["AWS_SECRET_ACCESS_KEY"] = ""
os.environ["AWS_REGION_NAME"] = ""

response = completion(
            model="sagemaker/jumpstart-dft-meta-textgeneration-llama-2-7b", 
            messages=[{ "content": "Hello, how are you?","role": "user"}],
            temperature=0.2,
            max_tokens=80,
            stream=True,
        )
for chunk in response:
    print(chunk)

LiteLLM Proxy Usage

Here's how to call Sagemaker with the LiteLLM Proxy Server

1. Setup config.yaml

model_list:
  - model_name: jumpstart-model
    litellm_params:
      model: sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614
      aws_access_key_id: os.environ/CUSTOM_AWS_ACCESS_KEY_ID
      aws_secret_access_key: os.environ/CUSTOM_AWS_SECRET_ACCESS_KEY
      aws_region_name: os.environ/CUSTOM_AWS_REGION_NAME

All possible auth params:

aws_access_key_id: Optional[str],
aws_secret_access_key: Optional[str],
aws_session_token: Optional[str],
aws_region_name: Optional[str],
aws_session_name: Optional[str],
aws_profile_name: Optional[str],
aws_role_name: Optional[str],
aws_web_identity_token: Optional[str],

2. Start the proxy

litellm --config /path/to/config.yaml

3. Test it

Curl Request

curl --location 'http://0.0.0.0:4000/chat/completions' \
--header 'Content-Type: application/json' \
--data ' {
      "model": "jumpstart-model",
      "messages": [
        {
          "role": "user",
          "content": "what llm are you"
        }
      ]
    }
'

OpenAI v1.0.0+

import openai
client = openai.OpenAI(
    api_key="anything",
    base_url="http://0.0.0.0:4000"
)

response = client.chat.completions.create(model="jumpstart-model", messages = [
    {
        "role": "user",
        "content": "this is a test request, write a short poem"
    }
])

print(response)

Langchain

from langchain.chat_models import ChatOpenAI
from langchain.prompts.chat import (
    ChatPromptTemplate,
    HumanMessagePromptTemplate,
    SystemMessagePromptTemplate,
)
from langchain.schema import HumanMessage, SystemMessage

chat = ChatOpenAI(
    openai_api_base="http://0.0.0.0:4000", # set openai_api_base to the LiteLLM Proxy
    model = "jumpstart-model",
    temperature=0.1
)

messages = [
    SystemMessage(
        content="You are a helpful assistant that im using to make a test request to."
    ),
    HumanMessage(
        content="test from litellm. tell me why it's amazing in 1 sentence"
    ),
]
response = chat(messages)

print(response)

Set temperature, top p, etc.

SDK

import os
from litellm import completion

os.environ["AWS_ACCESS_KEY_ID"] = ""
os.environ["AWS_SECRET_ACCESS_KEY"] = ""
os.environ["AWS_REGION_NAME"] = ""

response = completion(
  model="sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614",
  messages=[{ "content": "Hello, how are you?","role": "user"}],
  temperature=0.7,
  top_p=1
)

PROXY

Set on yaml

model_list:
  - model_name: jumpstart-model
    litellm_params:
      model: sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614
      temperature: <your-temp>
      top_p: <your-top-p>

Set on request

import openai
client = openai.OpenAI(
    api_key="anything",
    base_url="http://0.0.0.0:4000"
)

# request sent to model set on litellm proxy, `litellm --model`
response = client.chat.completions.create(model="jumpstart-model", messages = [
    {
        "role": "user",
        "content": "this is a test request, write a short poem"
    }
],
temperature=0.7,
top_p=1
)

print(response)

Allow setting temperature=0 for Sagemaker

By default when temperature=0 is sent in requests to LiteLLM, LiteLLM rounds up to temperature=0.1 since Sagemaker fails most requests when temperature=0

If you want to send temperature=0 for your model here's how to set it up (Since Sagemaker can host any kind of model, some models allow zero temperature)

SDK

import os
from litellm import completion

os.environ["AWS_ACCESS_KEY_ID"] = ""
os.environ["AWS_SECRET_ACCESS_KEY"] = ""
os.environ["AWS_REGION_NAME"] = ""

response = completion(
  model="sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614",
  messages=[{ "content": "Hello, how are you?","role": "user"}],
  temperature=0,
  aws_sagemaker_allow_zero_temp=True,
)

PROXY

Set aws_sagemaker_allow_zero_temp on yaml

model_list:
  - model_name: jumpstart-model
    litellm_params:
      model: sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614
      aws_sagemaker_allow_zero_temp: true

Set temperature=0 on request

import openai
client = openai.OpenAI(
    api_key="anything",
    base_url="http://0.0.0.0:4000"
)

# request sent to model set on litellm proxy, `litellm --model`
response = client.chat.completions.create(model="jumpstart-model", messages = [
    {
        "role": "user",
        "content": "this is a test request, write a short poem"
    }
],
temperature=0,
)

print(response)

Pass provider-specific params

If you pass a non-openai param to litellm, we'll assume it's provider-specific and send it as a kwarg in the request body. See more

SDK

import os
from litellm import completion

os.environ["AWS_ACCESS_KEY_ID"] = ""
os.environ["AWS_SECRET_ACCESS_KEY"] = ""
os.environ["AWS_REGION_NAME"] = ""

response = completion(
  model="sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614",
  messages=[{ "content": "Hello, how are you?","role": "user"}],
  top_k=1 # PROVIDER-SPECIFIC PARAM
)

PROXY

Set on yaml

model_list:
  - model_name: jumpstart-model
    litellm_params:
      model: sagemaker/jumpstart-dft-hf-textgeneration1-mp-20240815-185614
      top_k: 1 # PROVIDER-SPECIFIC PARAM

Set on request

import openai
client = openai.OpenAI(
    api_key="anything",
    base_url="http://0.0.0.0:4000"
)

# request sent to model set on litellm proxy, `litellm --model`
response = client.chat.completions.create(model="jumpstart-model", messages = [
    {
        "role": "user",
        "content": "this is a test request, write a short poem"
    }
],
temperature=0.7,
extra_body={
    "top_k": 1 # PROVIDER-SPECIFIC PARAM
}
)

print(response)

Passing Inference Component Name

If you have multiple models on an endpoint, you'll need to specify the individual model names, do this via model_id.

import os 
from litellm import completion

os.environ["AWS_ACCESS_KEY_ID"] = ""
os.environ["AWS_SECRET_ACCESS_KEY"] = ""
os.environ["AWS_REGION_NAME"] = ""

response = completion(
            model="sagemaker/<your-endpoint-name>", 
            model_id="<your-model-name",
            messages=[{ "content": "Hello, how are you?","role": "user"}],
            temperature=0.2,
            max_tokens=80
        )

Pin the inference component on a deployment

Since v1.109.0

When a proxy deployment or an SDK Router deployment sets model_id, LiteLLM will invoke that inference component and will ignore any model_id in the request body (unless the deployment allows the client to choose as per the below). A deployment can also pin the component via the X-Amzn-SageMaker-Inference-Component header in extra_headers and that pin will take precedence over the request in the same manner. On sagemaker_chat/ deployments LiteLLM will drop the client's own X-Amzn-SageMaker-Inference-Component header (regardless of case) if the deployment has pinned a component.

model_list:
  - model_name: my-llama
    litellm_params:
      model: sagemaker_chat/<your-endpoint-name>
      model_id: <llama-inference-component>
  - model_name: my-mistral
    litellm_params:
      model: sagemaker_chat/<your-endpoint-name>
      extra_headers:
        X-Amzn-SageMaker-Inference-Component: <mistral-inference-component>

This request still goes to <llama-inference-component>:

curl http://0.0.0.0:4000/v1/chat/completions \
  -H "Authorization: Bearer $LITELLM_API_KEY" \
  -H "Content-Type: application/json" \
  -d '{"model": "my-llama", "model_id": "<mistral-inference-component>", "messages": [{"role": "user", "content": "hi"}]}'

To allow the client to pick the component on a deployment, list model_id in its configurable_clientside_auth_params. A request that contains model_id is sent to that component and a request that doesn't contain it is sent to the deployment's own model_id. general_settings.allow_client_side_credentials: true doesn't allow the client to override a pinned component.

model_list:
  - model_name: my-endpoint
    litellm_params:
      model: sagemaker_chat/<your-endpoint-name>
      model_id: <default-inference-component>
      configurable_clientside_auth_params: ["model_id"]

A deployment that doesn't set either model_id or the header will still use the model_id from the request. Prior to 1.109.0 a model_id in the request body would have replaced the deployment's so a client could send a pinned deployment to a different component on the same endpoint.

Passing credentials as parameters - Completion()

Pass AWS credentials as parameters to litellm.completion

import os 
from litellm import completion

response = completion(
            model="sagemaker/jumpstart-dft-meta-textgeneration-llama-2-7b",
            messages=[{ "content": "Hello, how are you?","role": "user"}],
            aws_access_key_id="",
            aws_secret_access_key="",
            aws_region_name="",
)

Applying Prompt Templates

To apply the correct prompt template for your sagemaker deployment, pass in it's hf model name as well.

import os 
from litellm import completion

os.environ["AWS_ACCESS_KEY_ID"] = ""
os.environ["AWS_SECRET_ACCESS_KEY"] = ""
os.environ["AWS_REGION_NAME"] = ""

response = completion(
            model="sagemaker/jumpstart-dft-meta-textgeneration-llama-2-7b", 
            messages=messages,
            temperature=0.2,
            max_tokens=80,
            hf_model_name="meta-llama/Llama-2-7b",
        )

You can also pass in your own custom prompt template

Sagemaker Messages API

Use route sagemaker_chat/* to route to Sagemaker Messages API

model: sagemaker_chat/<your-endpoint-name>

SDK

import os
import litellm
from litellm import completion

litellm.set_verbose = True # SEE RAW REQUEST

os.environ["AWS_ACCESS_KEY_ID"] = ""
os.environ["AWS_SECRET_ACCESS_KEY"] = ""
os.environ["AWS_REGION_NAME"] = ""

response = completion(
            model="sagemaker_chat/<your-endpoint-name>", 
            messages=[{ "content": "Hello, how are you?","role": "user"}],
            temperature=0.2,
            max_tokens=80
        )

PROXY

1. Setup config.yaml

model_list:
  - model_name: "sagemaker-model"
    litellm_params:
      model: "sagemaker_chat/jumpstart-dft-hf-textgeneration1-mp-20240815-185614"
      aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID
      aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY
      aws_region_name: os.environ/AWS_REGION_NAME

2. Start the proxy

litellm --config /path/to/config.yaml

3. Test it

curl --location 'http://0.0.0.0:4000/chat/completions' \
--header 'Content-Type: application/json' \
--data ' {
      "model": "sagemaker-model",
      "messages": [
        {
          "role": "user",
          "content": "what llm are you"
        }
      ]
    }
'

** See OpenAI SDK/Langchain/Llamaindex/etc. examples**

Completion Models

tip

We support ALL Sagemaker models, just set model=sagemaker/<any-model-on-sagemaker> as a prefix when sending litellm requests

Here's an example of using a sagemaker model with LiteLLM

Model Name	Function Call
Your Custom Huggingface Model	completion(model='sagemaker/<your-deployment-name>', messages=messages)
Meta Llama 2 7B	completion(model='sagemaker/jumpstart-dft-meta-textgeneration-llama-2-7b', messages=messages)
Meta Llama 2 7B (Chat/Fine-tuned)	completion(model='sagemaker/jumpstart-dft-meta-textgeneration-llama-2-7b-f', messages=messages)
Meta Llama 2 13B	completion(model='sagemaker/jumpstart-dft-meta-textgeneration-llama-2-13b', messages=messages)
Meta Llama 2 13B (Chat/Fine-tuned)	completion(model='sagemaker/jumpstart-dft-meta-textgeneration-llama-2-13b-f', messages=messages)
Meta Llama 2 70B	completion(model='sagemaker/jumpstart-dft-meta-textgeneration-llama-2-70b', messages=messages)
Meta Llama 2 70B (Chat/Fine-tuned)	completion(model='sagemaker/jumpstart-dft-meta-textgeneration-llama-2-70b-b-f', messages=messages)

Embedding Models

LiteLLM supports all Sagemaker Jumpstart Huggingface Embedding models. Here's how to call it:

from litellm import completion

os.environ["AWS_ACCESS_KEY_ID"] = ""
os.environ["AWS_SECRET_ACCESS_KEY"] = ""
os.environ["AWS_REGION_NAME"] = ""

response = litellm.embedding(model="sagemaker/<your-deployment-name>", input=["good morning from litellm", "this is another item"])
print(f"response: {response}")

Embeddings on inference components

If the endpoint hosts inference components, pass the component to call via model_id, same as for completions. LiteLLM sends it as the InferenceComponentName of the invoke call and keeps it out of the request body.

response = litellm.embedding(
    model="sagemaker/<your-endpoint-name>",
    model_id="<your-inference-component-name>",
    input=["good morning from litellm"],
)

OpenAI-compatible embedding endpoints

By default LiteLLM sends embedding requests in the Hugging Face shape ({"inputs": [...]}) and reads a list of vectors back. If the container behind your endpoint serves the OpenAI embeddings API instead ({"input": [...]} in, {"data": [{"embedding": [...]}], "usage": {...}} out), add the openai/ route in front of the endpoint name. The request then goes out in the OpenAI shape with any extra_body fields merged in, and data[].embedding and usage are read from the response.

SDK

response = litellm.embedding(
    model="sagemaker/openai/<your-endpoint-name>",
    model_id="<your-inference-component-name>",  # only for endpoints with inference components
    extra_body={"model": "<model name the container expects>"},  # only if the container needs one
    input=["good morning from litellm"],
)

PROXY

model_list:
  - model_name: my-embeddings
    litellm_params:
      model: sagemaker/openai/<your-endpoint-name>
      model_id: <your-inference-component-name>
      extra_body:
        model: <model name the container expects>
      aws_region_name: us-east-1
    model_info:
      mode: embedding

Provider-specific keys in litellm_params are sent in the request body as-is, so keep only the ones your container accepts.

Reject chat calls on embedding endpoints

Since v1.109.0

An embedding container answers a chat request after the round trip to SageMaker with its own error, or with an empty assistant message when it treats the prompt as text to embed. To have the proxy reject chat calls before they leave it, list the endpoints the deployment serves and turn on enforcement:

    model_info:
      mode: embedding
      supported_endpoints: ["/v1/embeddings"]
      enforce_supported_endpoints: true

Chat-style requests to that model then go to a configured fallback, or get a 400 that names the endpoint and the supported list, and /v1/embeddings keeps working. See Reject chat calls on non-chat models for the routes covered and how fallbacks are handled.

Nova Models on SageMaker

LiteLLM supports Amazon Nova models (Nova Micro, Nova Lite, Nova 2 Lite) deployed on SageMaker Inference real-time endpoints. These custom/fine-tuned Nova models use an OpenAI-compatible API format.

Reference: AWS Blog - Amazon SageMaker Inference for Custom Amazon Nova Models

Usage

Use the sagemaker_nova/ prefix with your SageMaker endpoint name:

import litellm
import os

os.environ["AWS_ACCESS_KEY_ID"] = ""
os.environ["AWS_SECRET_ACCESS_KEY"] = ""
os.environ["AWS_REGION_NAME"] = "us-east-1"

# Basic chat completion
response = litellm.completion(
    model="sagemaker_nova/my-nova-endpoint",
    messages=[{"role": "user", "content": "Hello, how are you?"}],
    temperature=0.7,
    max_tokens=512,
)
print(response.choices[0].message.content)

Streaming

response = litellm.completion(
    model="sagemaker_nova/my-nova-endpoint",
    messages=[{"role": "user", "content": "Write a short poem"}],
    stream=True,
    stream_options={"include_usage": True},
)
for chunk in response:
    if chunk.choices and chunk.choices[0].delta.content:
        print(chunk.choices[0].delta.content, end="")

Multimodal (Images)

Nova models on SageMaker support image inputs using base64 data URIs:

response = litellm.completion(
    model="sagemaker_nova/my-nova-endpoint",
    messages=[
        {
            "role": "user",
            "content": [
                {"type": "text", "text": "What's in this image?"},
                {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,..."}}
            ]
        }
    ],
)

Proxy Config

model_list:
  - model_name: nova-micro
    litellm_params:
      model: sagemaker_nova/my-nova-micro-endpoint
      aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID
      aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY
      aws_region_name: us-east-1

Supported Parameters

All standard OpenAI parameters are supported, plus these Nova-specific parameters:

Parameter	Type	Description
top_k	integer	Limits token selection to top K most likely tokens
reasoning_effort	"low" | "high"	Reasoning effort level (Nova 2 Lite custom models only)
allowed_token_ids	array[int]	Restrict output to specified token IDs
truncate_prompt_tokens	integer	Truncate prompt to N tokens if it exceeds limit

response = litellm.completion(
    model="sagemaker_nova/my-nova-endpoint",
    messages=[{"role": "user", "content": "Think step by step: what is 2+2?"}],
    top_k=40,
    reasoning_effort="low",
    logprobs=True,
    top_logprobs=2,
)