SageMaker Real-Time Inference from MCP Server Tools

Invoking a SageMaker endpoint from an MCP tool handler

An MCP tool that wraps a SageMaker classifier endpoint — the handler serializes the input, invokes the endpoint, and deserializes the response:

import boto3
import json
from botocore.exceptions import ClientError

sagemaker_runtime = boto3.client("sagemaker-runtime", region_name="us-east-1")
ENDPOINT_NAME = "mcp-intent-classifier-prod"

def classify_intent(user_input: str) -> dict:
    """MCP tool handler: classify user intent using SageMaker endpoint."""
    payload = {
        "inputs": user_input,
        "parameters": {
            "max_length": 128,
            "return_all_scores": True,
        }
    }

    try:
        response = sagemaker_runtime.invoke_endpoint(
            EndpointName=ENDPOINT_NAME,
            ContentType="application/json",
            Accept="application/json",
            Body=json.dumps(payload),
        )
        # Body is a StreamingBody — must read() to get bytes
        result_bytes = response["Body"].read()
        result = json.loads(result_bytes)
        return {"intent": result[0]["label"], "score": result[0]["score"]}

    except sagemaker_runtime.exceptions.ModelError as e:
        # Container returned 4xx or 5xx — model-side error
        # e.response["Error"]["Code"] is the HTTP status from the container
        original_status = e.response.get("OriginalStatusCode", "unknown")
        original_message = e.response.get("OriginalMessage", str(e))
        raise RuntimeError(
            f"Model returned error {original_status}: {original_message}"
        ) from e

    except sagemaker_runtime.exceptions.InternalFailure as e:
        # SageMaker infrastructure error — retry is appropriate
        raise RuntimeError(f"SageMaker infrastructure error: {e}") from e

    except sagemaker_runtime.exceptions.ServiceUnavailable as e:
        # Endpoint is starting up or overloaded — retry with backoff
        raise RuntimeError(f"SageMaker endpoint unavailable: {e}") from e

The StreamingBody returned in response["Body"] must be consumed with .read() before the connection is released. Not reading it can cause connection pool exhaustion in long-running MCP servers. For large responses, use .read(amt) in a loop or .iter_chunks().

Three distinct exceptions to handle per call:

• ModelError — the model container itself returned a non-2xx HTTP response. Check OriginalStatusCode and OriginalMessage for the container's actual error. Common causes: malformed input shape, unsupported content type, input exceeds model's max sequence length.
• InternalFailure — SageMaker infrastructure failure. Retry with exponential backoff is correct here; the model itself may be healthy.
• ServiceUnavailable — endpoint is initializing, restarting after a failure, or scaling. Back off and retry; do not surface as a permanent tool error.

Content types and payload serialization

The ContentType and Accept headers must match what the model container's serving stack expects. Mismatches cause ModelError with status 415 or malformed output:

# JSON lines — used by Hugging Face TGI containers and many custom containers
response = sagemaker_runtime.invoke_endpoint(
    EndpointName=ENDPOINT_NAME,
    ContentType="application/json",
    Accept="application/json",
    Body=json.dumps({"inputs": text, "parameters": {"max_new_tokens": 256}}),
)

# CSV — used by built-in algorithms (XGBoost, Linear Learner)
csv_payload = ",".join(str(v) for v in feature_vector)
response = sagemaker_runtime.invoke_endpoint(
    EndpointName=ENDPOINT_NAME,
    ContentType="text/csv",
    Accept="text/csv",
    Body=csv_payload,
)

# Raw binary — image classification with built-in image models
with open("input.jpg", "rb") as f:
    image_bytes = f.read()
response = sagemaker_runtime.invoke_endpoint(
    EndpointName=ENDPOINT_NAME,
    ContentType="image/jpeg",
    Accept="application/json",
    Body=image_bytes,
)

# Tensor serialization — PyTorch endpoints using torchserve
import io, torch
tensor = torch.tensor(feature_vector, dtype=torch.float32)
buffer = io.BytesIO()
torch.save(tensor, buffer)
response = sagemaker_runtime.invoke_endpoint(
    EndpointName=ENDPOINT_NAME,
    ContentType="application/x-pytorch",
    Accept="application/json",
    Body=buffer.getvalue(),
)

For Hugging Face Inference Toolkit containers, the serving stack automatically handles JSON payloads with the inputs key. For TGI (Text Generation Inference) containers, additional parameters like temperature, top_p, and repetition_penalty go inside the parameters dict. For custom containers, the serialization contract is defined by the container's /invocations HTTP handler.

Endpoint auto-scaling for MCP traffic bursts

SageMaker endpoint instances don't scale automatically — you must configure Application Auto Scaling. The standard metric is SageMakerVariantInvocationsPerInstance:

import boto3

autoscaling = boto3.client("application-autoscaling", region_name="us-east-1")

resource_id = f"endpoint/{ENDPOINT_NAME}/variant/AllTraffic"

# Register the endpoint variant as a scalable target
autoscaling.register_scalable_target(
    ServiceNamespace="sagemaker",
    ResourceId=resource_id,
    ScalableDimension="sagemaker:variant:DesiredInstanceCount",
    MinCapacity=1,
    MaxCapacity=8,
)

# Target tracking: maintain ≤70 invocations per instance per minute
autoscaling.put_scaling_policy(
    PolicyName="mcp-endpoint-invocations-policy",
    ServiceNamespace="sagemaker",
    ResourceId=resource_id,
    ScalableDimension="sagemaker:variant:DesiredInstanceCount",
    PolicyType="TargetTrackingScaling",
    TargetTrackingScalingPolicyConfiguration={
        "TargetValue": 70.0,
        "PredefinedMetricSpecification": {
            "PredefinedMetricType": "SageMakerVariantInvocationsPerInstance",
        },
        "ScaleInCooldown": 300,   # 5 minutes — avoid thrashing
        "ScaleOutCooldown": 60,   # 1 minute — scale out fast for MCP bursts
    },
)

# For ML-latency-sensitive tools: scale on ModelLatency P99 instead
autoscaling.put_scaling_policy(
    PolicyName="mcp-endpoint-latency-policy",
    ServiceNamespace="sagemaker",
    ResourceId=resource_id,
    ScalableDimension="sagemaker:variant:DesiredInstanceCount",
    PolicyType="TargetTrackingScaling",
    TargetTrackingScalingPolicyConfiguration={
        "TargetValue": 500.0,   # target P99 latency ≤500ms
        "CustomizedMetricSpecification": {
            "MetricName": "ModelLatency",
            "Namespace": "/aws/sagemaker/Endpoints",
            "Dimensions": [
                {"Name": "EndpointName", "Value": ENDPOINT_NAME},
                {"Name": "VariantName", "Value": "AllTraffic"},
            ],
            "Statistic": "p99",
            "Unit": "Microseconds",
        },
        "ScaleInCooldown": 300,
        "ScaleOutCooldown": 60,
    },
)

Scale-in cooldown should be longer than scale-out cooldown — a new instance takes 2-5 minutes to initialize the model container, so scaling in too aggressively creates latency spikes when traffic returns. MinCapacity=1 ensures the endpoint always has at least one instance ready; setting MinCapacity=0 causes a cold start of 2-5 minutes on the first MCP tool call after a quiet period.

Multi-model endpoints for cost-efficient MCP tool variants

A Multi-Model Endpoint (MME) hosts multiple model artifacts behind one endpoint URL. SageMaker dynamically loads model artifacts from S3 on first invocation and caches them in instance memory. Use MME when an MCP server needs to call many model variants (per-tenant fine-tuned models, language-specific models) but traffic per model is too low to justify dedicated instances:

# MME invoke — add TargetModel parameter
response = sagemaker_runtime.invoke_endpoint(
    EndpointName="mcp-multi-model-endpoint",
    ContentType="application/json",
    Accept="application/json",
    TargetModel="tenant-acme-finetuned-v2.tar.gz",  # S3 key relative to ModelDataUrl prefix
    Body=json.dumps({"inputs": user_input}),
)

# List currently loaded models on the endpoint
sagemaker = boto3.client("sagemaker")
response = sagemaker.describe_endpoint(EndpointName="mcp-multi-model-endpoint")
# Check EndpointConfigName, then DescribeEndpointConfig for ModelDataUrl prefix

# The first invocation of a TargetModel triggers a download from S3 + load into memory
# Subsequent invocations use the cached model — typical first-invocation overhead: 10-60s
# To pre-warm: invoke with a dummy payload before serving real MCP traffic
def prewarm_model(target_model: str) -> None:
    try:
        sagemaker_runtime.invoke_endpoint(
            EndpointName="mcp-multi-model-endpoint",
            ContentType="application/json",
            TargetModel=target_model,
            Body=json.dumps({"inputs": "warmup"}),
        )
    except Exception:
        pass  # pre-warm failure is non-critical

MME constraints to know: all models on an MME must use the same container image and framework. Models are evicted from memory (LRU) when the instance runs out of memory — evicted models are reloaded from S3 on next invocation, incurring the same cold-start penalty. For MCP servers with strict latency SLAs, pre-warm the MME models that will be called in the session before the first user interaction.

Endpoint variants for A/B testing and canary rollouts

An endpoint can have multiple variants, each pointing to a different model version with configurable traffic weights. Use this for zero-downtime model upgrades in an MCP production environment:

sagemaker = boto3.client("sagemaker")

# Create endpoint config with two variants: 90% stable, 10% canary
sagemaker.create_endpoint_config(
    EndpointConfigName="mcp-classifier-ab-config",
    ProductionVariants=[
        {
            "VariantName": "Stable",
            "ModelName": "mcp-intent-classifier-v3",
            "InstanceType": "ml.m5.large",
            "InitialInstanceCount": 2,
            "InitialVariantWeight": 9,   # 90% traffic
        },
        {
            "VariantName": "Canary",
            "ModelName": "mcp-intent-classifier-v4",
            "InstanceType": "ml.m5.large",
            "InitialInstanceCount": 1,
            "InitialVariantWeight": 1,   # 10% traffic
        },
    ],
)

# Force a specific variant — useful in MCP A/B test logging
response = sagemaker_runtime.invoke_endpoint(
    EndpointName=ENDPOINT_NAME,
    ContentType="application/json",
    TargetVariant="Canary",    # override traffic routing
    Body=json.dumps(payload),
)
# response["InvokedProductionVariant"] tells you which variant handled the request

# Shift traffic when canary looks healthy: update variant weights
sagemaker.update_endpoint_weights_and_capacities(
    EndpointName=ENDPOINT_NAME,
    DesiredWeightsAndCapacities=[
        {"VariantName": "Stable", "DesiredWeight": 0},
        {"VariantName": "Canary", "DesiredWeight": 1},
    ],
)
# After confirming full traffic on Canary: update endpoint config to point Stable → v4
# and remove Canary variant to consolidate

InvokedProductionVariant in the response allows the MCP tool handler to log which model version produced each result — essential for offline evaluation and feedback loops. Combine variant logging with MCP tool result quality metrics to build a data flywheel that informs the next model version.