/async-inference

name: async-inference
description: Implement SageMaker async inference with S3-based I/O and polling. Use for long-running inference (>60s), large payloads, or batch processing workloads.

This skill guides implementation of SageMaker async inference, ideal for workloads that exceed real-time inference limits (60-second timeout, 6MB payload). Async inference stores inputs/outputs in S3 and supports polling or SNS notifications.

When to Use Async Inference

| Scenario | Inference Type | |----------|---------------| | Response < 60 seconds | Real-time | | Response 1-15 minutes | **Async** | | Payload > 6MB | **Async** | | Batch processing | **Async** or Batch Transform | | Cost-sensitive (scale to zero) | **Async** with auto-scaling |

Architecture

┌─────────┐    ┌──────────────┐    ┌─────────────────┐
│ Lambda  │───▶│ S3 (input)   │───▶│ SageMaker Async │
└─────────┘    └──────────────┘    │    Endpoint     │
                                   └────────┬────────┘
                                            │
              ┌──────────────┐              │
              │ S3 (output)  │◀─────────────┘
              │   or         │
              │ S3 (failure) │
              └──────┬───────┘
                     │
              ┌──────▼───────┐
              │ Poll / SNS   │
              └──────────────┘

Implementation

1. CDK Endpoint Configuration

import * as sagemaker from 'aws-cdk-lib/aws-sagemaker';

const asyncConfig = new sagemaker.CfnEndpointConfig(this, 'AsyncConfig', {
  endpointConfigName: 'my-async-endpoint-config',
  productionVariants: [{
    modelName: model.attrModelName,
    variantName: 'AllTraffic',
    instanceType: 'ml.g5.2xlarge',
    initialInstanceCount: 1,
  }],
  asyncInferenceConfig: {
    outputConfig: {
      s3OutputPath: `s3://${bucket.bucketName}/output/`,
      s3FailurePath: `s3://${bucket.bucketName}/failure/`,
      // Optional: SNS notification
      notificationConfig: {
        successTopic: successTopic.topicArn,
        errorTopic: errorTopic.topicArn,
      },
    },
    clientConfig: {
      maxConcurrentInvocationsPerInstance: 4,
    },
  },
});

2. TypeScript Client Implementation

import {
  SageMakerRuntimeClient,
  InvokeEndpointAsyncCommand
} from '@aws-sdk/client-sagemaker-runtime';
import {
  S3Client,
  GetObjectCommand,
  PutObjectCommand,
  ListObjectsV2Command
} from '@aws-sdk/client-s3';
import { v4 as uuid } from 'uuid';

interface AsyncInferenceResult<T> {
  success: boolean;
  data?: T;
  error?: string;
  inferenceId: string;
}

export class SageMakerAsyncClient {
  private sagemaker: SageMakerRuntimeClient;
  private s3: S3Client;
  private bucket: string;
  private endpointName: string;

  constructor(config: { bucket: string; endpointName: string; region: string }) {
    this.sagemaker = new SageMakerRuntimeClient({ region: config.region });
    this.s3 = new S3Client({ region: config.region });
    this.bucket = config.bucket;
    this.endpointName = config.endpointName;
  }

  async invoke<T>(payload: object): Promise<AsyncInferenceResult<T>> {
    const inferenceId = uuid();
    const inputKey = `input/${inferenceId}/request.json`;

    // 1. Upload input to S3
    await this.s3.send(new PutObjectCommand({
      Bucket: this.bucket,
      Key: inputKey,
      Body: JSON.stringify(payload),
      ContentType: 'application/json',
    }));

    // 2. Invoke async endpoint
    const response = await this.sagemaker.send(new InvokeEndpointAsyncCommand({
      EndpointName: this.endpointName,
      InputLocation: `s3://${this.bucket}/${inputKey}`,
      InvocationTimeoutSeconds: 3600,  // 1 hour max
      RequestTTLSeconds: 21600,        // 6 hour TTL
    }));

    // Validate OutputLocation is present
    const outputLocation = response.OutputLocation;
    if (!outputLocation) {
      return {
        success: false,
        error: 'OutputLocation not returned from async invocation',
        inferenceId,
      };
    }

    // 3. Poll for result
    return this.pollForResult<T>(inferenceId, outputLocation);
  }

  private async pollForResult<T>(
    inferenceId: string,
    outputLocation: string,
    maxWaitMs: number = 1200000,  // 20 minutes
    initialPollIntervalMs: number = 5000,
    maxPollIntervalMs: number = 30000
  ): Promise<AsyncInferenceResult<T>> {
    const startTime = Date.now();
    const outputKey = outputLocation.replace(`s3://${this.bucket}/`, '');
    // Extract output filename to construct failure path
    const outputFileName = outputKey.split('/').pop()!;
    const failureKey = `failure/${outputFileName}`;
    let pollIntervalMs = initialPollIntervalMs;

    while (Date.now() - startTime < maxWaitMs) {
      // Check for success
      try {
        const result = await this.s3.send(new GetObjectCommand({
          Bucket: this.bucket,
          Key: outputKey,
        }));

        if (!result.Body) {
          return {
            success: false,
            error: 'Empty response body from S3',
            inferenceId,
          };
        }

        const body = await result.Body.transformToString();

        // Parse JSON with explicit error handling
        let parsed: T;
        try {
          parsed = JSON.parse(body) as T;
        } catch (parseError: any) {
          return {
            success: false,
            error: `Failed to parse inference result JSON: ${parseError?.message ?? 'Unknown parse error'}`,
            inferenceId,
          };
        }

        return {
          success: true,
          data: parsed,
          inferenceId,
        };
      } catch (e: any) {
        // Only continue polling for missing objects
        if (e.name !== 'NoSuchKey' && e?.$metadata?.httpStatusCode !== 404) {
          throw e;
        }
      }

      // Check for failure
      try {
        const failures = await this.listPrefix(failureKey);
        if (failures.length > 0) {
          const errorBody = await this.getObject(fai