endpoint-two-models.yml

AWSTemplateFormatVersion: 2010-09-09

Parameters:
  ModelName:
    Description: First model name 
    Type: String
  ModelDataUrl:
    Description: Location of first model artefact
    Type: String
  ModelName2:
    Description: Second model name
    Type: String
  ModelDataUrl2:
    Description: Location of second model artefact
    Type: String
  TrainingImage:
    Description: The container used to train the models
    Type: String
  InstanceType:
    Description: Instance type
    Type: String
    Default: ml.t2.xlarge
  InstanceCount:
    Description: Instance count
    Type: String
    Default: 1
  RoleArn:
    Description: Execution Role ARN
    Type: String

Resources:
  Model:
    Type: "AWS::SageMaker::Model"
    Properties: 
      Containers: 
        - 
          Image: !Ref TrainingImage
          ModelDataUrl: !Ref ModelDataUrl
      ExecutionRoleArn: !Ref RoleArn
      ModelName: !Ref ModelName

  Model2:
    Type: "AWS::SageMaker::Model"
    Properties: 
      Containers: 
        - 
          Image: !Ref TrainingImage
          ModelDataUrl: !Ref ModelDataUrl2
      ExecutionRoleArn: !Ref RoleArn
      ModelName: !Ref ModelName2

  Endpoint:
    Type: "AWS::SageMaker::Endpoint"
    Properties: 
      EndpointConfigName: !GetAtt EndpointConfig.EndpointConfigName

  EndpointConfig:
    Type: "AWS::SageMaker::EndpointConfig"
    Properties:
      ProductionVariants: 
        -
          ModelName: !GetAtt Model.ModelName
          VariantName: variant-1
          InitialInstanceCount: !Ref InstanceCount
          InstanceType: !Ref InstanceType
          InitialVariantWeight: 1.0
          # The first model will get 2/3 of traffic
        -
          ModelName: !GetAtt Model2.ModelName
          VariantName: variant-2
          InitialInstanceCount: !Ref InstanceCount
          InstanceType: !Ref InstanceType
          InitialVariantWeight: 0.5
          # The second model will get 1/3 of traffic

Outputs:
  EndpointId:
    Value: !Ref Endpoint
  EndpointName:
    Value: !GetAtt Endpoint.EndpointName