【夏休みの自由研究リレー】 Snowflakeの外部関数からSageMaker AIの推論エンドポイントにデプロイしたChronos-2に時系列予測をリクエストする仕組みを作ってみた

"""Snowflake external function proxy for a JumpStart Chronos-2 endpoint."""

from __future__ import annotations

import base64
import hashlib
import json
import logging
import os
from typing import Any

import boto3

logger = logging.getLogger()
logger.setLevel(logging.INFO)

SAGEMAKER_ENDPOINT_NAME = os.environ.get("SAGEMAKER_ENDPOINT_NAME", "").strip()
runtime = boto3.client("sagemaker-runtime")

def _snowflake_error(status_code: int, message: str) -> dict[str, Any]:
    body = json.dumps({"data": [], "error": message})
    return {
        "statusCode": status_code,
        "body": body,
        "headers": {"Content-Type": "application/json"},
    }

def _snowflake_success(rows: list[list[Any]]) -> dict[str, Any]:
    body = json.dumps({"data": rows})
    md5digest = hashlib.md5(body.encode("utf-8")).digest()
    return {
        "statusCode": 200,
        "body": body,
        "headers": {
            "Content-Type": "application/json",
            "Content-MD5": base64.b64encode(md5digest).decode("ascii"),
        },
    }

def _parse_event_body(event: dict[str, Any]) -> dict[str, Any]:
    raw_body = event.get("body")
    if raw_body is None:
        raise ValueError("Request body is missing")

    if event.get("isBase64Encoded"):
        raw_body = base64.b64decode(raw_body).decode("utf-8")

    return json.loads(raw_body)

def handler(event: dict[str, Any], context: Any) -> dict[str, Any]:
    try:
        if not SAGEMAKER_ENDPOINT_NAME:
            raise ValueError("SAGEMAKER_ENDPOINT_NAME is not set")

        rows = _parse_event_body(event)["data"]
        prediction_length, freq = int(rows[0][4]), str(rows[0][5])

        inputs = []
        for row in rows:
            if int(row[4]) != prediction_length or str(row[5]) != freq:
                raise ValueError("prediction_length and freq must match within a batch")
            inputs.append(
                {
                    "item_id": str(row[1]),
                    "start": str(row[2]),
                    "target": [float(value) for value in row[3]],
                }
            )

        payload = {
            "inputs": inputs,
            "parameters": {
                "prediction_length": prediction_length,
                "freq": freq,
            },
        }
        response = runtime.invoke_endpoint(
            EndpointName=SAGEMAKER_ENDPOINT_NAME,
            ContentType="application/json",
            Accept="application/json",
            Body=json.dumps(payload).encode("utf-8"),
        )
        forecasts = json.loads(response["Body"].read())["predictions"]

        output_rows = []
        for row, input_data, forecast in zip(rows, inputs, forecasts, strict=True):
            forecast.setdefault("item_id", input_data["item_id"])
            output_rows.append([int(row[0]), forecast])

        return _snowflake_success(output_rows)
    except Exception as exc:
        logger.exception("Chronos-2 external function failed")
        return _snowflake_error(400, str(exc))

AWSTemplateFormatVersion: "2010-09-09"
Description: API Gateway and Lambda proxy for Snowflake external functions calling Chronos-2 on SageMaker.

Parameters:
  SageMakerEndpointName:
    Type: String
    Description: Name of the Chronos-2 realtime endpoint deployed from SageMaker JumpStart.

  LambdaExecutionRoleName:
    Type: String
    AllowedPattern: "^[a-zA-Z0-9+=,.@_-]+$"
    Description: IAM role name used by the Lambda function.

  LambdaCodeS3Bucket:
    Type: String
    Description: S3 bucket containing the Lambda deployment package.

  LambdaCodeS3Key:
    Type: String
    Default: lambda/chronos2-external-function.zip
    Description: S3 object key for the Lambda deployment package.

  apiGatewayStageName:
    Type: String
    AllowedPattern: "^[-a-z0-9]+$"
    Default: ext-func-stage
    Description: API deployment stage.

  lambdaName:
    Type: String
    AllowedPattern: "^[a-zA-Z0-9]+[-a-zA-Z0-9-]+[-a-zA-Z0-9]+$"
    Default: chronos2-ext-func-lambda
    Description: Lambda function name.

  apiGatewayType:
    Type: String
    Default: REGIONAL
    AllowedValues:
      - REGIONAL
      - PRIVATE
    Description: API Gateway endpoint type.

  apiGatewayName:
    Type: String
    AllowedPattern: "^[a-zA-Z0-9]+[-a-zA-Z0-9-]+[-a-zA-Z0-9]+$"
    Default: chronos2-ext-func-api
    Description: API Gateway instance name.

  apiGatewayIAMRoleName:
    Type: String
    AllowedPattern: "^[a-zA-Z0-9]+[-a-zA-Z0-9-]+[-a-zA-Z0-9]+$"
    Description: IAM role assumed by Snowflake for API Gateway invocation.

  sourceVpcId:
    Type: String
    Default: ""
    Description: Snowflake VPC ID. Required only for PRIVATE API Gateway.

Conditions:
  shouldCreateRegionalGateway:
    !Equals [!Ref apiGatewayType, REGIONAL]

Resources:
  apiIAMRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: !Ref apiGatewayIAMRoleName
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              AWS: !Sub "arn:aws:iam::${AWS::AccountId}:root"
            Action: sts:AssumeRole

  lambdaExecutionIAMRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: !Ref LambdaExecutionRoleName
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
      Policies:
        - PolicyName: InvokeChronos2Endpoint
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action:
                  - sagemaker:InvokeEndpoint
                Resource: !Sub "arn:aws:sagemaker:${AWS::Region}:${AWS::AccountId}:endpoint/${SageMakerEndpointName}"

  lambdaFunction:
    Type: AWS::Lambda::Function
    DependsOn:
      - lambdaExecutionIAMRole
    Properties:
      FunctionName: !Ref lambdaName
      Description: Snowflake external function proxy for Chronos-2 on SageMaker.
      Runtime: python3.12
      Handler: handler.handler
      Role: !GetAtt lambdaExecutionIAMRole.Arn
      Timeout: 29
      MemorySize: 2048
      Environment:
        Variables:
          SAGEMAKER_ENDPOINT_NAME: !Ref SageMakerEndpointName
      Code:
        S3Bucket: !Ref LambdaCodeS3Bucket
        S3Key: !Ref LambdaCodeS3Key

  apiGateway:
    Type: AWS::ApiGateway::RestApi
    DependsOn:
      - apiIAMRole
    Properties:
      Name: !Ref apiGatewayName
      Description: Snowflake external functions gateway for Chronos-2.
      Policy: !Sub
        - |
          {
            "Version": "2012-10-17",
            "Statement": [
              {
                "Effect": "Allow",
                "Principal": {
                  "AWS": "arn:aws:sts::${AWS::AccountId}:assumed-role/${apiGatewayIAMRoleName}/snowflake"
                },
                "Action": "execute-api:Invoke",
                "Resource": "${resourceArn}",
                "Condition": { ${vpcCondition} }
              }
            ]
          }
        - resourceArn: !Join ["", ["execute-api:/", "*"]]
          vpcCondition: !If
            - shouldCreateRegionalGateway
            - ""
            - !Sub '"StringEquals": { "aws:sourceVpc": "${sourceVpcId}" }'
      EndpointConfiguration:
        Types:
          - !Ref apiGatewayType

  apiResource:
    Type: AWS::ApiGateway::Resource
    Properties:
      RestApiId: !Ref apiGateway
      ParentId: !GetAtt
        - apiGateway
        - RootResourceId
      PathPart: chronos2

  apiGatewayRootMethod:
    Type: AWS::ApiGateway::Method
    Properties:
      AuthorizationType: AWS_IAM
      HttpMethod: POST
      Integration:
        IntegrationHttpMethod: POST
        Type: AWS_PROXY
        Uri: !Sub
          - arn:aws:apigateway:${AWS::Region}:lambda:path/2015-03-31/functions/${lambdaArn}/invocations
          - lambdaArn: !GetAtt lambdaFunction.Arn
      ResourceId: !Ref apiResource
      RestApiId: !Ref apiGateway

  apiGatewayDeployment:
    Type: AWS::ApiGateway::Deployment
    DependsOn:
      - apiGatewayRootMethod
    Properties:
      RestApiId: !Ref apiGateway
      StageName: !Ref apiGatewayStageName

  lambdaApiGatewayInvoke:
    Type: AWS::Lambda::Permission
    Properties:
      Action: lambda:InvokeFunction
      FunctionName: !GetAtt lambdaFunction.Arn
      Principal: apigateway.amazonaws.com
      SourceArn: !Sub arn:aws:execute-api:${AWS::Region}:${AWS::AccountId}:${apiGateway}/*/*/*

Outputs:
  ResourceInvocationUrl:
    Description: Resource invocation URL for CREATE EXTERNAL FUNCTION.
    Value: !Sub https://${apiGateway}.execute-api.${AWS::Region}.amazonaws.com/${apiGatewayStageName}/chronos2

  ApiGatewayIAMRoleArn:
    Description: IAM role ARN for Snowflake CREATE API INTEGRATION.
    Value: !GetAtt apiIAMRole.Arn

  LambdaExecutionRoleArn:
    Description: IAM role ARN used by the Lambda function.
    Value: !GetAtt lambdaExecutionIAMRole.Arn

  LambdaFunctionName:
    Description: Lambda function name.
    Value: !Ref lambdaFunction