Triton Inferene ServerをつかったSageMaker AIを構築する

概要

SageMaker AIをCloudFormationで構築し、Triton Inferene Serverを使ってデプロイした際の話をします。
⚠️ 今回の話はAWSのSageMaker AIの話が中心で、モデルについては言及しません。

SageMaker AIに必要なリソース

SageMaker AIに必要なリソースは以下の3つになります。

• SageMaker Endpoint: 実際にリクエストを処理するインスタンス + ALB等のリソース
• SageMaker Endpoint Configuration: SageMaker Endpointを構成するうえでの設定リソース
• SageMaker Model: SageMakerで利用するモデル

これにプラスして

• S3パケット
• ECR

等も必要です。

ref: https://pages.awscloud.com/rs/112-TZM-766/images/AWS-Black-Belt_2022_Amazon-SageMaker-Inference-Part-3_1014_v1.pdf

リソース構築

まずS3，ECRについては以下で構築します。

Resources:
  S3BucketSageMakerModels:
    Type: AWS::S3::Bucket
    DeletionPolicy: Retain
    UpdateReplacePolicy: Retain
    Properties:
      BucketName: sagemaker-models
      PublicAccessBlockConfiguration:
        BlockPublicAcls: true
        BlockPublicPolicy: true
        IgnorePublicAcls: true
        RestrictPublicBuckets: true
      BucketEncryption:
        ServerSideEncryptionConfiguration:
          - ServerSideEncryptionByDefault:
              SSEAlgorithm: AES256
      VersioningConfiguration:
        Status: Enabled
      Tags:
        - Key: CostEnv
          Value: !Ref CostEnv
        - Key: CostService
          Value: !Ref CostService
  ECRRepositorySageMaker:
    Type: AWS::ECR::Repository
    Properties:
      RepositoryName: "sagemaker-tritonserver"
      ImageTagMutability: !If
        - IsProduction
        - IMMUTABLE
        - !Ref AWS::NoValue
      Tags:
        - Key: CostEnv
          Value: !Ref CostEnv
        - Key: CostService
          Value: !Ref CostService

IAMロールは以下で構築します。

• SageMakerの権限
• S3アクセスへの権限
• LogGroupへの書き込み権限

今回は行いませんが、SageMakerはVpcConfigでVPCを指定することも出来ます。

  IAMRoleSageMaker:
    Type: AWS::IAM::Role
    Properties:
      RoleName: "sagemaker-role"
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: sagemaker.amazonaws.com
            Action: sts:AssumeRole
      Policies:
        - PolicyName: SageMakerAccess
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - "sagemaker:*"
                  - "cloudwatch:PutMetricData"
                  - "ecr:GetAuthorizationToken"
                  - "ecr:BatchCheckLayerAvailability"
                  - "ecr:GetDownloadUrlForLayer"
                  - "ecr:BatchGetImage"
                Resource: "*"
        - PolicyName: S3ModelDataRead
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - "s3:GetObject"
                  - "s3:ListBucket"
                Resource:
                  - "arn:aws:s3:::sagemaker-models"
                  - "arn:aws:s3:::sagemaker-models/*"
        - PolicyName: VPCAccess
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - "ec2:DescribeVpcEndpoints"
                  - "ec2:DescribeDhcpOptions"
                  - "ec2:DescribeVpcs"
                  - "ec2:DescribeSubnets"
                  - "ec2:DescribeSecurityGroups"
                  - "ec2:DescribeNetworkInterfaces"
                  - "ec2:DeleteNetworkInterfacePermission"
                  - "ec2:DeleteNetworkInterface"
                  - "ec2:CreateNetworkInterfacePermission"
                  - "ec2:CreateNetworkInterface"
                Resource: "*"
        - PolicyName: CloudWatchLogsWrite
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - "logs:CreateLogGroup"
                  - "logs:CreateLogStream"
                  - "logs:PutLogEvents"
                Resource: "*"
      Tags:
        - Key: CostEnv
          Value: !Ref CostEnv
        - Key: CostService
          Value: !Ref CostService

モデルは事前にアップロードをしておきます。

aws s3 cp model_7.tar.gz s3://sagemaker-models/artifacts/model_7.tar.gz

またDockerイメージは事前にAWSの941853720454アカウントからPullしてきて、自アカウントのECRにpushをしておきます。

$ aws ecr get-login-password --region ap-northeast-1 | docker login --username AWS --password-stdin 941853720454.dkr.ecr.ap-northeast-1.amazonaws.com

$ docker pull 941853720454.dkr.ecr.ap-northeast-1.amazonaws.com/sagemaker-tritonserver:23.10-py3

$ docker tag 941853720454.dkr.ecr.ap-northeast-1.amazonaws.com/sagemaker-tritonserver:23.10-py3 717397055837.dkr.ecr.ap-northeast-1.amazonaws.com/sagemaker-tritonserver:23.10-py3

$ docker push 111111111.dkr.ecr.ap-northeast-1.amazonaws.com/sagemaker-tritonserver:23.10-py3

また941853720454アカウントから持ってくるdockerイメージはAMDのみなので、インスタンス指定する場合もAMD or Intelのインスタンスである必要があります。

❯ docker inspect 941853720454.dkr.ecr.ap-northeast-1.amazonaws.com/sagemaker-tritonserver:23.08-py3 | jq -r '.[0].Architecture'
amd64

次にSageMaker Model, SageMakerEndpointConfig、SageMakerEndpointについても構築します。

  SageMakerModel:
    Type: AWS::SageMaker::Model
    Properties:
      ExecutionRoleArn: IAMRoleSageMaker
      PrimaryContainer:
        Image: !Sub "${AWS::AccountId}.dkr.ecr.ap-northeast-1.amazonaws.com/sagemaker-tritonserver:23.10-py3"
        ModelDataUrl: !Sub "s3://sagemaker-models/artifacts/${ModelFileName}"

  SageMakerEndpointConfig:
    Type: AWS::SageMaker::EndpointConfig
    Properties:
      ProductionVariants:
        - ModelName: !GetAtt SageMakerModel.ModelName
          VariantName: AllTraffic
          InstanceType: ml.r5.xlarge
          InitialInstanceCount: 1
          InitialVariantWeight: 1.0
          ContainerStartupHealthCheckTimeoutInSeconds: 60

  SageMakerEndpoint:
    Type: AWS::SageMaker::Endpoint
    Properties:
      EndpointName: sagemaker-endpoint
      EndpointConfigName: !GetAtt SageMakerEndpointConfig.EndpointConfigName
      Tags:
        - Key: CostEnv
          Value: !Ref CostEnv
        - Key: CostService
          Value: !Ref CostService

これでapplyすればリソース構築が始まります。
またSageMakerを構成する各リソースは設定をどれかでもいじって再applyすると、他リソースの再リリースも行ってくれる点も便利で気にいっています。