SageMaker Async Inference from MCP Server Tools

Configuring an async inference endpoint

Async inference is configured at the endpoint config level via AsyncInferenceConfig. The same model artifact and container are used — only the serving mode changes:

import boto3

sagemaker = boto3.client("sagemaker", region_name="us-east-1")
s3_bucket = "mcp-sagemaker-async-io"

# Create endpoint config with async inference enabled
sagemaker.create_endpoint_config(
    EndpointConfigName="mcp-doc-processor-async-config",
    ProductionVariants=[
        {
            "VariantName": "AllTraffic",
            "ModelName": "mcp-doc-processor-v2",
            "InstanceType": "ml.g5.2xlarge",
            "InitialInstanceCount": 1,
        }
    ],
    AsyncInferenceConfig={
        "OutputConfig": {
            "S3OutputPath": f"s3://{s3_bucket}/outputs/",
            "NotificationConfig": {
                # SNS topic ARN — receives message on success AND failure
                "SuccessTopic": "arn:aws:sns:us-east-1:123456789012:mcp-inference-success",
                "ErrorTopic": "arn:aws:sns:us-east-1:123456789012:mcp-inference-error",
            },
            # KMS key for encrypting output in S3
            "KmsKeyId": "arn:aws:kms:us-east-1:123456789012:key/...",
        },
        "ClientConfig": {
            # Max time SageMaker will wait for model response (1 to 3600 seconds)
            "MaxConcurrentInvocationsPerInstance": 4,
        },
    },
)

MaxConcurrentInvocationsPerInstance controls how many async requests a single instance processes simultaneously. Higher values improve throughput but increase per-request latency if the model is compute-bound. For GPU-backed large language models, set this to 1-2; for CPU-bound document processing, 4-8 is typical.

Submitting async inference requests from MCP tool handlers

The MCP tool handler uploads the input payload to S3, then calls invoke_endpoint_async:

import boto3
import json
import uuid
import time
from datetime import datetime

s3 = boto3.client("s3", region_name="us-east-1")
sagemaker_runtime = boto3.client("sagemaker-runtime", region_name="us-east-1")

ENDPOINT_NAME = "mcp-doc-processor-async"
INPUT_BUCKET = "mcp-sagemaker-async-io"
OUTPUT_BUCKET = "mcp-sagemaker-async-io"

def process_document_async(document_text: str, job_id: str | None = None) -> dict:
    """Submit a document processing job to SageMaker async inference."""
    job_id = job_id or str(uuid.uuid4())
    input_key = f"inputs/{job_id}.json"
    output_prefix = f"s3://{OUTPUT_BUCKET}/outputs/"

    # Upload input to S3
    payload = {
        "document": document_text,
        "options": {"extract_entities": True, "summarize": True},
    }
    s3.put_object(
        Bucket=INPUT_BUCKET,
        Key=input_key,
        Body=json.dumps(payload),
        ContentType="application/json",
    )

    # Submit async invocation — returns immediately
    response = sagemaker_runtime.invoke_endpoint_async(
        EndpointName=ENDPOINT_NAME,
        ContentType="application/json",
        InputLocation=f"s3://{INPUT_BUCKET}/{input_key}",
        InvocationTimeoutSecs=300,   # fail after 5 minutes if no response
        RequestTTLSecs=600,          # drop from queue if not started within 10 minutes
    )

    # OutputLocation is the full S3 path where the result will appear
    # e.g. s3://mcp-sagemaker-async-io/outputs/{inference-id}.out
    return {
        "inference_id": response["InferenceId"],
        "output_location": response["OutputLocation"],
        "submitted_at": datetime.utcnow().isoformat() + "Z",
    }

RequestTTLSecs is the maximum time a request can sit in the async queue waiting for an available instance. If the queue is full and RequestTTLSecs expires, the request is dropped with no output — only an SNS error notification (if configured). Set it to at least 2× InvocationTimeoutSecs to avoid queue-expiry races when the endpoint is under load.

Polling for async inference results

The simplest integration pattern for MCP tools: the handler polls S3 for the output file, blocking the tool call until the result is ready or a timeout expires:

def poll_async_result(
    output_location: str,
    failure_prefix: str | None = None,
    poll_interval: float = 2.0,
    max_wait: float = 300.0,
) -> dict:
    """
    Poll S3 until async inference result appears or timeout.
    output_location: full S3 URI from invoke_endpoint_async response
    failure_prefix:  S3 URI prefix where SageMaker writes failure details
    """
    # Parse S3 URI
    # output_location format: s3://bucket/outputs/{InferenceId}.out
    assert output_location.startswith("s3://")
    parts = output_location[5:].split("/", 1)
    bucket, key = parts[0], parts[1]

    # Derive failure path: same key with .out replaced by error info
    # SageMaker writes to {output_prefix}{InferenceId}.out or {InferenceId}
    base_key = key.rsplit(".", 1)[0] if "." in key else key
    failure_key = f"{base_key}.failure"

    deadline = time.monotonic() + max_wait
    while time.monotonic() < deadline:
        # Check for successful output
        try:
            s3_client = boto3.client("s3")
            obj = s3_client.get_object(Bucket=bucket, Key=key)
            return json.loads(obj["Body"].read())
        except s3_client.exceptions.NoSuchKey:
            pass

        # Check for failure output
        try:
            fail_obj = s3_client.get_object(Bucket=bucket, Key=failure_key)
            error_detail = fail_obj["Body"].read().decode("utf-8")
            raise RuntimeError(f"Inference failed: {error_detail}")
        except s3_client.exceptions.NoSuchKey:
            pass

        time.sleep(poll_interval)
        poll_interval = min(poll_interval * 1.2, 10.0)  # gentle backoff, cap at 10s

    raise TimeoutError(
        f"Async inference did not complete within {max_wait}s. "
        f"Check {output_location} manually."
    )

The polling approach is appropriate when the MCP tool's response is needed synchronously within the same conversation turn. For very long jobs (>60s), consider returning the InferenceId as part of the tool response and letting the LLM call a second MCP tool (check_job_status) in a follow-up turn — this avoids holding an open MCP connection for minutes.

SNS-based completion notifications

For production MCP deployments, SNS notifications are more reliable than polling — they avoid S3 read costs and don't require a sleeping thread per outstanding request:

# SNS notification payload on SUCCESS:
# {
#   "eventVersion": "1.0",
#   "eventSource": "aws:SageMaker",
#   "awsRegion": "us-east-1",
#   "inferenceId": "abc123-...",
#   "endpointName": "mcp-doc-processor-async",
#   "outputLocation": "s3://mcp-sagemaker-async-io/outputs/abc123-.../output",
#   "contentType": "application/json",
# }

# SNS notification payload on FAILURE:
# {
#   "eventVersion": "1.0",
#   "eventSource": "aws:SageMaker",
#   "inferenceId": "abc123-...",
#   "failureCode": "ModelError",         # or "InvocationTimeoutExceeded"
#   "failureMessage": "Container returned 500",
#   "outputLocation": null,
#   "failurePath": "s3://mcp-sagemaker-async-io/outputs/abc123-.../failure",
# }

# Lambda function subscribed to the SNS success topic
def handle_inference_complete(event, context):
    for record in event["Records"]:
        sns_message = json.loads(record["Sns"]["Message"])
        inference_id = sns_message["inferenceId"]
        output_location = sns_message["outputLocation"]

        # Retrieve result from S3
        bucket, key = parse_s3_uri(output_location)
        s3 = boto3.client("s3")
        result = json.loads(s3.get_object(Bucket=bucket, Key=key)["Body"].read())

        # Store result in DynamoDB keyed by inference_id for MCP tool polling
        dynamodb = boto3.resource("dynamodb")
        table = dynamodb.Table("mcp-async-results")
        table.put_item(Item={
            "inference_id": inference_id,
            "result": result,
            "completed_at": datetime.utcnow().isoformat() + "Z",
            "ttl": int(time.time()) + 3600,  # expire after 1 hour
        })

# MCP tool: check_job_status
def check_job_status(inference_id: str) -> dict:
    dynamodb = boto3.resource("dynamodb")
    table = dynamodb.Table("mcp-async-results")
    item = table.get_item(Key={"inference_id": inference_id}).get("Item")
    if item is None:
        return {"status": "pending", "inference_id": inference_id}
    return {"status": "complete", "result": item["result"]}

The SNS → Lambda → DynamoDB pattern creates a results cache that MCP tool handlers can query without polling S3 directly. The ttl attribute in DynamoDB automatically cleans up old results, preventing unbounded table growth. Set the TTL to at least the maximum session duration for the agent using the MCP server.

Choosing between async, real-time, and batch transform

Three SageMaker inference modes serve different MCP tool workloads:

• Real-time inference — use when latency <30s is required and requests arrive continuously. Best for conversational MCP tools (intent classification, entity extraction, short-form generation). Cost: always-on instances, even during idle periods.
• Async inference — use when model response takes 30s-3600s, or when traffic is bursty and queuing is acceptable. Best for document processing, image/video generation, long-context summarization. Cost: instances scale to zero when queue is empty (with auto-scaling configured).
• Batch transform — use for offline processing of large datasets, not interactive MCP tools. SageMaker creates instances, processes S3 input files, writes output to S3, then terminates instances. No endpoint URL — triggered via create_transform_job. Best for nightly batch scoring or bulk document pre-processing that feeds a cache used by MCP tools.

Async inference auto-scaling with MinCapacity=0 allows the endpoint to scale to zero between requests, reducing cost significantly for low-traffic MCP tools. The trade-off: the first request after a scale-to-zero event waits 2-5 minutes for a new instance to start. Pre-scale to 1 instance before a known peak period by temporarily setting MinCapacity=1.