Update register script to send CloudWatch event

Author: brightsparc

deployment_pipeline/README.md

@@ -10,16 +10,18 @@ This deployment pipeline contains a few stages.
 1. **Source**: Pull the latest deployment configuration from AWS CodeCommit repository.
 1. **Build**: AWS CodeBuild job to create the AWS CloudFormation template for deploying the endpoint.
     - Query the Amazon SageMaker project to get the top approved models.
-    - Use the AWS CDK to create a CFN stack with multiple endpoint variants.
-    - Create a `register.json` file that contains the target SageMaker endpoint and A/B testing strategy.
-2. **Deploy**: Run the AWS CloudFormation stack to create/update the SageMaker endpoint.
-3. **Register**: Call the RegisterAPI for the endpoint to create/clear the A/B testing metrics.
+    - Use the AWS CDK to create a CFN stack to deploy multi-variant SageMaker Endpoint.
+2. **Deploy**: Run the AWS CloudFormation stack to create/update the SageMaker endpoint, tagged with properties based on configuration:
+    - `ab-testing:enabled` equals `true`
+    - `ab-testing:strategy` is one `WeightedSampling`, `EpslionGreedy`, `UCB1` or `ThompsonSampling`.
+    - `ab-testing:epsilon` is parameters for `EpslionGreedy` strategy, defaults to `0.1`.
+    - `ab-testing:warmup` the number of invocations to warmup with `WeightedSampling` strategy, defaults to `0`.
 
 ![\[AWS CodePipeline\]](../docs/ab-testing-pipeline-code-pipeline.png)
 
 ## Testing
 
-Once you have created a SageMaker Project, you can test the **Build** and **Register** stages locally by setting some environment variables, and running the commands found in the `buildspec` defined in the pipeline.
+Once you have created a SageMaker Project, you can test the **Build** stage and **Register** events locally by setting some environment variables.
 
 ### Build Stage
 
@@ -32,12 +34,11 @@ export STAGE_NAME="dev"
 cdk synth
 ```
 
-### Register Stage
+### Register
 
 Export the environment variable for the `REGISTER_LAMBDA` created as part of the `ab-testing-api` stack, then run `register.py` file.
 
 ```
 export REGISTER_LAMBDA="<<register_lambda>>"
 python register.py
 ```
-

deployment_pipeline/app.py

@@ -45,7 +45,7 @@
     # 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:strategy", 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)),
     ]

deployment_pipeline/infra/model_registry.py

@@ -59,7 +59,7 @@ def get_latest_approved_packages(
             # Return error if no packages found
             if len(model_packages) == 0 and creation_time_after is None:
                 error_message = (
-                    f"No latest packages found for: {model_package_group_name}"
+                    f"No approved packages found for: {model_package_group_name}"
                 )
                 logger.error(error_message)
                 raise Exception(error_message)
@@ -121,9 +121,7 @@ def get_versioned_approved_packages(
 
             # Return error if no packages found
             if len(model_packages) == 0:
-                error_message = (
-                    f"No versioned packages found for: {model_package_group_name}"
-                )
+                error_message = f"No approved packages found for: {model_package_group_name} and versions: {model_package_versions}"
                 logger.error(error_message)
                 raise Exception(error_message)
 

deployment_pipeline/register.py

@@ -24,7 +24,20 @@
 # Get the config and include with endpoint to register this model
 with open(f"{stage_name}-config.json", "r") as f:
     j = json.load(f)
-    event = json.dumps({"endpoint_name": endpoint_name, **j})
+    event = json.dumps({
+        'source': 'aws.sagemaker',
+        'detail-type': 'SageMaker Endpoint State Change',
+        'detail': {
+            'EndpointName': endpoint_name,
+            'EndpointStatus': 'IN_SERVICE',
+            'Tags': {
+                'ab-testing:enabled': 'true',
+                'ab-testing:strategy': j.get('strategy', 'ThompsonSampling'),
+                'ab-testing:epsilon': str(j.get('epsilon', 0.1)),
+                'ab-testing:warmup': str(j.get('warmup', 0)),
+            }
+        }
+    })
     response = lambda_client.invoke(
         FunctionName=register_lambda,
         InvocationType="RequestResponse",

docs/ab-testing-pipeline-code-pipeline.png

@@ -190,6 +190,9 @@ def __init__(
             targets=[targets.LambdaFunction(lambda_register)],
         )
 
+        # 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(
             actions=["cloudwatch:PutMetricData"], resources=["*"]

infra/api_stack.py

@@ -111,7 +111,7 @@ def lambda_handler(event, context):
         if endpoint_status == "DELETING":
             result, status_code = handle_delete(endpoint_name)
         elif endpoint_status == "IN_SERVICE":
-            strategy = endpoint_tags["ab-testing:stategy"]
+            strategy = endpoint_tags["ab-testing:strategy"]
             epsilon = float(endpoint_tags["ab-testing:epsilon"])
             warmup = int(endpoint_tags["ab-testing:warmup"])
             result, status_code = handle_register(