Update API stack to register SageMaker Endpoints from cloudwatch events, removed stage from CodePipeline deployment.

Author: brightsparc

README.md

@@ -146,7 +146,7 @@ Follow are a list of context values that are provided in the `cdk.json`, which c
 |---------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------|------------------------------------|
 | `api_name`                | The API Gateway Name                                                                                                                                            | "ab-testing"                       |
 | `stage_name`              | The stage namespace for resource and API Gateway path                                                                                                           | "dev"                              |
-| `endpoint_filter`         | A prefix to filter which Amazon SageMaker endpoints the API can invoke                                                                                          | "*"                                |
+| `endpoint_prefix`         | A prefix to filter which Amazon SageMaker endpoints the API can invoked.                                                                                        | ""                                 |
 | `api_lambda_memory`       | The [lambda memory](https://docs.aws.amazon.com/lambda/latest/dg/configuration-memory.html) allocation for API endpoint.                                        | 768                                |
 | `api_lambda_timeout`      | The lambda timeout for the API endpoint.                                                                                                                        | 10                                 |
 | `metrics_lambda_memory`   | The [lambda memory](https://docs.aws.amazon.com/lambda/latest/dg/configuration-memory.html) allocated for metrics processing Lambda                             | 768                                |
@@ -164,9 +164,7 @@ Run the following command to deploy the API and testing infrastructure, optional
 cdk deploy ab-testing-api
 ```
 
-This stack will ask you to confirm any changes, and output the `RegisterLambda` which you will provide to the MLOps Project, and the `ApiEndpoint` which you will provide to the A/B Testing sample notebook.
-
-Amazon SageMaker Studio projects will be granted access to invoke the Register Lambda, so if you are seeing errors running the above command ensure you have [Enable SageMaker project templates for Studio users](https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-projects-studio-updates.html).
+This stack will ask you to confirm any changes, and output the `ApiEndpoint` which you will provide to the A/B Testing sample notebook.
 
 ## Create the SageMaker MLOps Project Template
 
@@ -276,7 +274,6 @@ On the Create project page, SageMaker templates is chosen by default. This optio
   - The project name must have 32 characters or fewer.
 10. In the Project template parameters, provide the *Repository Name** you created previously eg:
   - For **StageName**, enter `dev` 
-  - For **RegisterLambda**, enter the `RegisterLambda` output from the `ab-testing-api` stack
   - For **CodeCommitSeedBucket**, enter the `CodeCommitSeedBucket` output from the `ab-testiing-service-catalog` stack
   - For **CodeCommitSeedKey**, enter the `CodeCommitSeedKey` output from the `ab-testiing-service-catalog` stack
 11. Choose Create project.

ab-testing-pipeline.yml

@@ -23,10 +23,6 @@ Parameters:
     Type: String
     Description: The optional s3 seed key
     MinLength: 1
-  RegisterLambda:
-    Type: String
-    Description: The AWS Lambda to invoke when registering this model
-    MinLength: 1
 Resources:
   CodeRepo:
     Type: AWS::CodeCommit::Repository
@@ -55,8 +51,6 @@ Resources:
         - Key: sagemaker:project-name
           Value:
             Ref: SageMakerProjectName
-    Metadata:
-      aws:cdk:path: ab-testing-pipeline/CodeRepo
   CdkBuild455F642E:
     Type: AWS::CodeBuild::Project
     Properties:
@@ -131,76 +125,6 @@ Resources:
             - Ref: SageMakerProjectName
             - -cdk-
             - Ref: StageName
-    Metadata:
-      aws:cdk:path: ab-testing-pipeline/CdkBuild/Resource
-  RegisterBuild8D2EA4DC:
-    Type: AWS::CodeBuild::Project
-    Properties:
-      Artifacts:
-        Type: CODEPIPELINE
-      Environment:
-        ComputeType: BUILD_GENERAL1_SMALL
-        EnvironmentVariables:
-          - Name: SAGEMAKER_PROJECT_NAME
-            Type: PLAINTEXT
-            Value:
-              Ref: SageMakerProjectName
-          - Name: STAGE_NAME
-            Type: PLAINTEXT
-            Value:
-              Ref: StageName
-          - Name: REGISTER_LAMBDA
-            Type: PLAINTEXT
-            Value:
-              Ref: RegisterLambda
-        Image: aws/codebuild/standard:1.0
-        ImagePullCredentialsType: CODEBUILD
-        PrivilegedMode: false
-        Type: LINUX_CONTAINER
-      ServiceRole:
-        Fn::Join:
-          - ""
-          - - "arn:"
-            - Ref: AWS::Partition
-            - ":iam::"
-            - Ref: AWS::AccountId
-            - :role/service-role/AmazonSageMakerServiceCatalogProductsUseRole
-      Source:
-        BuildSpec: >-
-          {
-            "version": "0.2",
-            "phases": {
-              "build": {
-                "commands": [
-                  "python register.py > output.txt"
-                ]
-              }
-            },
-            "artifacts": {
-              "files": [
-                "output.txt"
-              ]
-            },
-            "environment": {
-              "buildImage": {
-                "type": "LINUX_CONTAINER",
-                "defaultComputeType": "BUILD_GENERAL1_SMALL",
-                "imageId": "aws/codebuild/amazonlinux2-x86_64-standard:3.0",
-                "imagePullPrincipalType": "CODEBUILD"
-              }
-            }
-          }
-        Type: CODEPIPELINE
-      EncryptionKey: alias/aws/s3
-      Name:
-        Fn::Join:
-          - ""
-          - - sagemaker-
-            - Ref: SageMakerProjectName
-            - -register-
-            - Ref: StageName
-    Metadata:
-      aws:cdk:path: ab-testing-pipeline/RegisterBuild/Resource
   S3Artifact80610462:
     Type: AWS::S3::Bucket
     Properties:
@@ -215,8 +139,6 @@ Resources:
             - Ref: AWS::Region
     UpdateReplacePolicy: Delete
     DeletionPolicy: Delete
-    Metadata:
-      aws:cdk:path: ab-testing-pipeline/S3Artifact/Resource
   PipelineC660917D:
     Type: AWS::CodePipeline::Pipeline
     Properties:
@@ -292,22 +214,6 @@ Resources:
               Name: SageMaker_CFN_Deploy
               RunOrder: 1
           Name: Deploy
-        - Actions:
-            - ActionTypeId:
-                Category: Build
-                Owner: AWS
-                Provider: CodeBuild
-                Version: "1"
-              Configuration:
-                ProjectName:
-                  Ref: RegisterBuild8D2EA4DC
-              InputArtifacts:
-                - Name: Artifact_Source_CodeCommit_Source
-              Name: Register_Build
-              OutputArtifacts:
-                - Name: Artifact_Register_Register_Build
-              RunOrder: 1
-          Name: Register
       ArtifactStore:
         Location:
           Ref: S3Artifact80610462
@@ -319,8 +225,6 @@ Resources:
             - Ref: SageMakerProjectName
             - -pipeline-
             - Ref: StageName
-    Metadata:
-      aws:cdk:path: ab-testing-pipeline/Pipeline/Resource
   DeployRule0F8E909D:
     Type: AWS::Events::Rule
     Properties:
@@ -369,8 +273,6 @@ Resources:
                 - ":iam::"
                 - Ref: AWS::AccountId
                 - :role/service-role/AmazonSageMakerServiceCatalogProductsUseRole
-    Metadata:
-      aws:cdk:path: ab-testing-pipeline/DeployRule/Resource
   CodeRule663E3DC0:
     Type: AWS::Events::Rule
     Properties:
@@ -430,14 +332,10 @@ Resources:
                 - ":iam::"
                 - Ref: AWS::AccountId
                 - :role/service-role/AmazonSageMakerServiceCatalogProductsUseRole
-    Metadata:
-      aws:cdk:path: ab-testing-pipeline/CodeRule/Resource
   CDKMetadata:
     Type: AWS::CDK::Metadata
     Properties:
       Analytics: v2:deflate64:H4sIAAAAAAAAE0WN0YrCMBBFv8X3OCIVln0T/YFSvyBOR3ZskynJRCkh/742Lfh0L3M5Z47we4Lj7mzfcY/9cMgogSDf1OJgrg/f2mAdKQXTUZQUkMxVfNSQUJf9e334z9CzsvhiFl1m6yC3MjLOVVVbMSg9oTjHClUwSWSVMNfhnnjsPxRPNLKnNsiT1k9bLSY2kC8JB6rnta3WaaO+fAW3Xgy9yGuE3KWxLkuWUkw765/4QwM/0OyekXkfkld2BN2a//2ekvEmAQAA
-    Metadata:
-      aws:cdk:path: ab-testing-pipeline/CDKMetadata/Default
     Condition: CDKMetadataAvailable
 Conditions:
   CDKMetadataAvailable:

cdk.json

@@ -11,7 +11,7 @@
     "log_level": "INFO",
     "api_name": "ab-testing",
     "stage_name": "dev",
-    "endpoint_filter": "*",
+    "endpoint_prefix": "",
     "api_lambda_memory": 768,
     "api_lambda_timeout": 60,
     "metrics_lambda_memory": 768,

deployment_pipeline/app.py

@@ -42,6 +42,13 @@
 with open(f"{stage_name}-config.json", "r") as f:
     j = json.load(f)
     deployment_config = DeploymentConfig(**j)
+    # Append tags for ab-testing
+    tags += [
+        core.CfnTag(key="ab-testing:enabled", value="true"),
+        core.CfnTag(key="ab-testing:stategy", value=deployment_config.strategy),
+        core.CfnTag(key="ab-testing:epsilon", value=str(deployment_config.epsilon)),
+        core.CfnTag(key="ab-testing:warmup", value=str(deployment_config.warmup)),
+    ]
 
 sagemaker = SageMakerStack(
     app,

deployment_pipeline/infra/deployment_config.py

@@ -42,6 +42,7 @@ def __init__(
         instance_type: str = "ml.t2.medium",
         strategy: str = "ThompsonSampling",
         warmup: int = 0,
+        epsilon: float = 0.1,
     ):
         self.stage_name = stage_name
         # Provide either the challenger variant count, or specific champion/challenger config
@@ -74,4 +75,5 @@ def __init__(
             self.challenger_variant_config = None
         self.strategy = strategy
         self.warmup = warmup
+        self.epsilon = epsilon
         super().__init__(instance_count, instance_type)

deployment_pipeline/prod-config.json

@@ -1,7 +1,8 @@
 {
     "stage_name": "prod",
-    "strategy": "ThompsonSampling",
+    "strategy": "EpsilonGreedy",
     "warmup": 100,
+    "epsilon": 0.1,
     "instance_count": 2,
     "instance_type": "ml.c5.large",
     "champion_variant_config": {

docs/ab-testing-pipeline-architecture.png

@@ -2,6 +2,8 @@
     core,
     aws_apigateway,
     aws_iam,
+    aws_events as events,
+    aws_events_targets as targets,
     aws_logs,
     aws_lambda,
     aws_dynamodb,
@@ -24,7 +26,7 @@ def __init__(
         log_level = self.node.try_get_context("log_level")
         api_name = self.node.try_get_context("api_name")
         stage_name = self.node.try_get_context("stage_name")
-        endpoint_filter = self.node.try_get_context("endpoint_filter")
+        endpoint_prefix = self.node.try_get_context("endpoint_prefix")
         api_lambda_memory = self.node.try_get_context("api_lambda_memory")
         api_lambda_timeout = self.node.try_get_context("api_lambda_timeout")
         metrics_lambda_memory = self.node.try_get_context("metrics_lambda_memory")
@@ -114,9 +116,9 @@ def __init__(
                     "sagemaker:InvokeEndpoint",
                 ],
                 resources=[
-                    "arn:aws:sagemaker:{}:{}:endpoint/{}".format(
-                        self.region, self.account, endpoint_filter
-                    ),
+                    "arn:aws:sagemaker:{}:{}:endpoint/{}*".format(
+                        self.region, self.account, endpoint_prefix
+                    )
                 ],
             )
         )
@@ -145,6 +147,7 @@ def __init__(
                 "DELIVERY_STREAM_NAME": delivery_stream_name,
                 "DELIVERY_SYNC": "true" if delivery_sync else "false",
                 "LOG_LEVEL": log_level,
+                "ENDPOINT_PREFIX": endpoint_prefix,
             },
             layers=[xray_layer],
             tracing=aws_lambda.Tracing.ACTIVE,
@@ -160,23 +163,32 @@ def __init__(
                     "sagemaker:DescribeEndpoint",
                 ],
                 resources=[
-                    "arn:aws:sagemaker:{}:{}:endpoint/{}".format(
-                        self.region, self.account, endpoint_filter
-                    ),
+                    "arn:aws:sagemaker:{}:{}:endpoint/{}*".format(
+                        self.region, self.account, endpoint_prefix
+                    )
                 ],
             )
         )
 
-        # Grant permissions to the service catalog use role
-        service_catalog_role = aws_iam.Role.from_role_arn(
+        # Add endpoint event rule to register endpoints that are created or updated.
+        # Note CDK is unable to filter on resource prefixes, so we will need to filter on this within the RegisterLambda function.
+        # see: https://docs.aws.amazon.com/eventbridge/latest/userguide/eb-event-patterns-content-based-filtering.html#filtering-prefix-matching
+        endpoint_rule = events.Rule(
             self,
-            "RegisterRole",
-            f"arn:{self.partition}:iam::{self.account}:role/service-role/AmazonSageMakerServiceCatalogProductsUseRole",
+            "EndpointRule",
+            rule_name=f"sagemaker-{api_name}-endpoint",
+            description="Rule to register a SageMaker Endpoint when it is created or updated.",
+            event_pattern=events.EventPattern(
+                source=["aws.sagemaker"],
+                detail_type=[
+                    "SageMaker Endpoint State Change",
+                ],
+                detail={
+                    "EndpointStatus": ["IN_SERVICE", "DELETING"],
+                },
+            ),
+            targets=[targets.LambdaFunction(lambda_register)],
         )
-        lambda_register.grant_invoke(service_catalog_role)
-
-        # Return the register lambda function as output
-        core.CfnOutput(self, "RegisterLambda", value=lambda_register.function_name)
 
         # Get cloudwatch put metrics policy ()
         cloudwatch_metric_policy = aws_iam.PolicyStatement(