Canary Deployments for A/B Testing AI Models.

Question

Pulumi · Accepted Answer

Canary deployments are a pattern in deployment strategies where a new version of a service is rolled out to a small subset of users before being rolled out to the entire infrastructure. This allows for testing the new version in a live environment with real traffic, while minimizing the impact of any potential issues to a small segment of the user base.

In the context of A/B Testing AI models, a canary deployment could be used to direct a portion of the traffic to a model variant to compare its performance against the currently deployed model. This approach is useful for assessing the effectiveness of a new model under real-world conditions without fully replacing the existing one.

Let's say you are using AWS for your infrastructure, and you want to set up a canary deployment to A/B test two different AI models using AWS services. On AWS, this can be achieved using services like AWS Lambda for running the AI models and AWS CodeDeploy for managing the deployment strategies.

Below is a Pulumi program that sets up a canary deployment for A/B testing two versions of an AWS Lambda function, which could represent two different AI models. We will use the `aws.codedeploy.Application` and `aws.codedeploy.DeploymentGroup` to set up the deployment with the canary configuration.

```python
import pulumi
import pulumi_aws as aws

# Create a Lambda Role to execute Lambda Functions
lambda_role = aws.iam.Role("lambdaRole",
    assume_role_policy=aws.iam.get_policy_document(statements=[{
        "actions": ["sts:AssumeRole"],
        "principals": [{
            "identifiers": ["lambda.amazonaws.com"],
            "type": "Service",
        }],
    }]).json
)

# Attach the AWSLambdaBasicExecutionRole policy to the role
attach_exec_policy = aws.iam.RolePolicyAttachment("lambdaBasicExecRole",
    role=lambda_role.name,
    policy_arn="arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole"
)

# Define the first version of the AI model Lambda function
ai_model_v1 = aws.lambda_.Function("aiModelV1",
    role=lambda_role.arn,
    description="Version 1 of AI Model",
    handler="index.handler",
    runtime="python3.8",
    code=pulumi.FileArchive("./path_to_code/v1/"),
    publish=True,  # This ensures that a new version is created on each deploy
)

# Define the second version of the AI model Lambda function (simulating the new A/B test model)
ai_model_v2 = aws.lambda_.Function("aiModelV2",
    role=lambda_role.arn,
    description="Version 2 of AI Model",
    handler="index.handler",
    runtime="python3.8",
    code=pulumi.FileArchive("./path_to_code/v2/"),
    publish=True,
)

# Create the CodeDeploy application for the lambda functions
cd_app = aws.codedeploy.Application("codeDeployApp",
    compute_platform="Lambda"
)

# Set up the deployment group for the canary deployment
# This example assumes 10% of traffic will be shifted initially, and the rest after 1 day
deployment_group = aws.codedeploy.DeploymentGroup("codeDeployDeploymentGroup",
    service_role_arn=lambda_role.arn,
    application_name=cd_app.name,
    deployment_config_name="CodeDeployDefault.LambdaAllAtOnce",
    deployment_style={
        "deployment_option": "WITH_TRAFFIC_CONTROL",
        "deployment_type": "BLUE_GREEN",
    },
    lambda_settings={
        "alarms": [{
            "name": "Errors",
            "enabled": True,
        }],
        "name": ai_model_v1.name,
    },
    auto_rollback_configuration={
        "enabled": True,
        "events": ["DEPLOYMENT_FAILURE", "DEPLOYMENT_STOP_ON_ALARM"],
    },
    trigger_configurations=[{
        "trigger_events": ["DeploymentFailure"],
        "trigger_name": "DeploymentFailure",
        "trigger_target_arn": "arn:aws:sns:REGION:ACCOUNT_ID:SNS_TOPIC_NAME",
    }],
    blue_green_deployment_config={
        "deployment_ready_option": {
            "action_on_timeout": "CONTINUE_DEPLOYMENT",
            "wait_time_in_minutes": 0,
        },
        "terminate_blue_instances_on_deployment_success": {
            "action": "TERMINATE",
            "termination_wait_time_in_minutes": 0,
        },
    },
)

pulumi.export("aiModelV1LambdaArn", ai_model_v1.arn)
pulumi.export("aiModelV2LambdaArn", ai_model_v2.arn)
pulumi.export("codeDeployApp", cd_app.name)
pulumi.export("deploymentGroup", deployment_group.name)
```

This program sets up two AWS Lambda functions, one representing the original AI model and the other representing the new model. We then create a CodeDeploy application and deployment group with a configuration that supports canary deployment.

You will notice that we have specified a `deployment_config_name` which defines how the traffic is routed between the old and new versions. In a real scenario, you would want to customize the percentages and timers according to the granularity of control you want over the traffic shift.

The `blue_green_deployment_config` is responsible for what happens once the deployment is deemed successful — for instance, terminating instances running the old code.

The `trigger_configurations` is an example setup that integrates with Amazon SNS to notify in case of deployment failure. You would need to replace `REGION`, `ACCOUNT_ID`, and `SNS_TOPIC_NAME` with the appropriate values specific to your environment.

Finally, we export the ARNs (Amazon Resource Names) of the Lambda functions and the names of the CodeDeploy Application and Deployment Group.

Make sure to replace `./path_to_code/v1/` and `./path_to_code/v2/` with the actual paths to the zipped source code of your two AI model versions. And update the trigger configuration with the correct SNS topic ARN for notifications.

This setup allows a portion of the user traffic to be directed to the new model, enabling A/B testing in a production environment with reduced risk. If the new version performs well, it can gradually receive all the traffic; otherwise, the system will roll back to the previous version automatically.