Skip to content

Develop a GenAI Solution with CDK CI/CD Workbench Part 2

In this section, we will continue developing the GenAI solution using the Workbench. We will create an AWS ECS Fargate task to run a Streamlit application that will host a Python web application to generate images using Amazon Bedrock.

Enhance the Solution

Step 1: Get your IP address

Before proceeding, we need to obtain your IP address to ensure that access to the application is restricted to only you.

  1. Open the following URL: https://checkip.amazonaws.com/.

    This will display your current IP address.

  2. Note down your IP address, as you will need it later to configure access restrictions for the application.

Step 2: Create the Fargate Task Definition

We’ll start by enhancing the solution with an ECS Fargate task definition to run the Streamlit application.

  1. Open the file lib/demo-stack.ts and add the following code after the ECS Cluster definition:
const taskRole = new iam.Role(this, 'TaskRole', {
    assumedBy: new iam.ServicePrincipal('ecs-tasks.amazonaws.com'),
    inlinePolicies: {
        BedrockInvokeModel: new iam.PolicyDocument({
        statements: [
            new iam.PolicyStatement({
            actions: ['bedrock:InvokeModel'],
            resources: [`arn:aws:bedrock:${cdk.Aws.REGION}::foundation-model/amazon.titan-image-generator-v1`],
            }),
        ],
        }),
    },
});

// Define a Fargate task definition
const fargateTaskDefinition = new ecs.FargateTaskDefinition(this, 'StreamlitAppTaskDef', {
    taskRole,
});

// Add container to task definition
const container = fargateTaskDefinition.addContainer('StreamlitContainer', {
    image: ecs.ContainerImage.fromAsset(path.resolve(__dirname, 'image_prompts'), { networkMode: ecrAsset.NetworkMode.custom('sagemaker') }),
    memoryLimitMiB: 512,
    cpu: 256,
    logging: new ecs.AwsLogDriver({ streamPrefix: 'streamlit-app' }),
});

// Expose the container on port 8501 (Streamlit default port)
container.addPortMappings({
    containerPort: 8501,
    protocol: ecs.Protocol.TCP,
});

// Create security group for ALB
const albSecurityGroup = new ec2.SecurityGroup(this, 'ALBSecurityGroup', {
    vpc: cluster.vpc,
    allowAllOutbound: true,
    description: 'Security group for ALB with restricted ingress',
});

for (const ip of allowedIPs) {
    albSecurityGroup.addIngressRule(
        ec2.Peer.ipv4(ip),
        ec2.Port.tcp(80),
        `Allow HTTP from ${ip}`,
    );
}

// Create Fargate service in ECS
const loadBalancer = new ecsPatterns.ApplicationLoadBalancedFargateService(this, 'StreamlitFargateService', {
    cluster: cluster, // Required
    taskDefinition: fargateTaskDefinition,
    desiredCount: 1,
    publicLoadBalancer: true,
    openListener: false,
});

loadBalancer.loadBalancer.addSecurityGroup(albSecurityGroup);

nag.NagSuppressions.addResourceSuppressions(
    loadBalancer,
    [
        {
        id: 'AwsSolutions-ELB2',
        reason: 'ELB access logs',
        },
    ],
    true,
);

nag.NagSuppressions.addResourceSuppressions(fargateTaskDefinition, [{
    id: 'AwsSolutions-IAM5',
    reason: 'Allow * permissions.',
    appliesTo: ['Resource::*']
}], true);
  1. Add the following import statements at the top of the file:
import * as iam from 'aws-cdk-lib/aws-iam';
import * as ec2 from 'aws-cdk-lib/aws-ec2';
import * as ecrAsset from 'aws-cdk-lib/aws-ecr-assets';
import * as ecsPatterns from 'aws-cdk-lib/aws-ecs-patterns';
import * as path from 'path';
import * as nag from 'cdk-nag';
  1. Add the following code to the allowedIPs array in the lib/demo-stack.ts file:
const allowedIPs = ['<your-ip-address>/32'];

Congratulations! You've successfully defined the Fargate task to run the Streamlit app.

Step 3: Create the Streamlit Application

Let’s create the Streamlit application that will interact with Amazon Bedrock for generating images.

Step 3.1: Create the Docker Setup

  1. Create a new directory called image_prompts in the lib folder.

  2. Create a new file called Dockerfile in the image_prompts folder.

  3. Add the following content to the Dockerfile:

FROM python:3.12-slim

# Set the working directory
WORKDIR /app

# Install pipenv to manage dependencies
RUN pip install pipenv

# Copy the Pipfile and Pipfile.lock first to leverage Docker cache
COPY Pipfile Pipfile.lock ./

# Install project dependencies from Pipfile
RUN pipenv install --system --deploy --ignore-pipfile

# Copy the rest of the application code
COPY . .

# Expose the port on which Streamlit runs
EXPOSE 8501

# Command to run Streamlit app
CMD ["streamlit", "run", "image_prompts_app.py", "--server.port", "8501", "--server.address", "0.0.0.0"]

Well done! The Docker setup for the Streamlit app is complete.

Step 3.2: Define Dependencies and Application Code

  1. Create a new file called Pipfile in the image_prompts folder.
  2. Add the following content to the Pipfile:
[[source]]
url = "https://pypi.python.org/simple"
verify_ssl = true
name = "pypi"

[packages]
streamlit = "*"
botocore = "*"
boto3 = "*"

[requires]
python_version = "3"
  1. Create a new file called Pipfile.lock in the image_prompts folder. Open a terminal and run the following command to generate the Pipfile.lock file:
cd lib/image_prompts && pipenv lock && cd ../..

Nice work! The dependencies for the Streamlit app have been set up.

Step 3.3: Add the Streamlit Application Logic

  1. Create a new file called image_prompts_app.py in the image_prompts folder.
  2. Add the following content to image_prompts_app.py:
import streamlit as st
import image_prompts_lib as glib


st.set_page_config(layout="wide", page_title="Image Generation")

st.title("Image Generation")

col1, col2 = st.columns(2)


with col1:
    st.subheader("Image parameters")

    prompt_text = st.text_area("What you want to see in the image:", height=100, help="The prompt text")
    negative_prompt = st.text_input("What shoud not be in the image:", help="The negative prompt")

    generate_button = st.button("Generate", type="primary")


with col2:
    st.subheader("Result")

    if generate_button:
        with st.spinner("Drawing..."):
            generated_image = glib.get_image_from_model(
                prompt_content=prompt_text, 
                negative_prompt=negative_prompt,
            )

        st.image(generated_image)

Awesome! The Streamlit application is ready to generate images.

Step 3.4: Create the Helper Library

  1. Create a new file called image_prompts_lib.py in the image_prompts folder.
  2. Add the following content to image_prompts_lib.py:
import boto3
import json
import base64
from io import BytesIO
from random import randint

#get the stringified request body for the InvokeModel API call
def get_titan_image_generation_request_body(prompt, negative_prompt=None):

    body = { #create the JSON payload to pass to the InvokeModel API
        "taskType": "TEXT_IMAGE",
        "textToImageParams": {
            "text": prompt,
        },
        "imageGenerationConfig": {
            "numberOfImages": 1,  # Number of images to generate
            "quality": "premium",
            "height": 512,
            "width": 512,
            "cfgScale": 8.0,
            "seed": randint(0, 100000),  #nosec Use a random seed
        },
    }

    if negative_prompt:
        body['textToImageParams']['negativeText'] = negative_prompt

    return json.dumps(body)


#get a BytesIO object from the Titan Image Generator response
def get_titan_response_image(response):

    response = json.loads(response.get('body').read())

    images = response.get('images')

    image_data = base64.b64decode(images[0])

    return BytesIO(image_data)


#generate an image using Amazon Titan Image Generator
def get_image_from_model(prompt_content, negative_prompt=None):
    session = boto3.Session()

    bedrock = session.client(service_name='bedrock-runtime') #creates a Bedrock client

    body = get_titan_image_generation_request_body(prompt_content, negative_prompt=negative_prompt)

    response = bedrock.invoke_model(body=body, modelId="amazon.titan-image-generator-v1", contentType="application/json", accept="application/json")

    output = get_titan_response_image(response)

    return output

Fantastic! The helper library is in place to interact with AWS Bedrock.

Step 4: Deploy the Workbench

Run the following command to deploy the Workbench environment:

npm run workbench deploy -- --all

Well done! The Workbench has been successfully deployed.

Step 5: Verify the Deployment

The deployment process may take a few minutes. Once it's complete, the LoadBalancer URL will be displayed in the terminal.

  1. Open the LoadBalancer URL in your browser to verify the deployment. The Streamlit application should be accessible.

Congratulations! Your application is up and running.

Introduce Feature Flag

Now, we need to introduce a feature flag to control whether the Streamlit application should be deployed. This allows us to work on the feature without immediately deploying it through the pipeline.

Step 6: Add the Feature Flag to Control the Streamlit App

  1. Open the lib/demo-stack.ts file and wrap the Streamlit app logic inside a feature flag check. The feature flag allows us to selectively enable or disable the deployment of the Streamlit app:

    if (cdk.FeatureFlags.of(this).isEnabled('feature-streamlit')) {
    wrapper.logger.info('Feature Streamlit is enabled');

    const taskRole = new iam.Role(this, 'TaskRole', {
    // ... Streamlit app definition
    });

    nag.NagSuppressions.addResourceSuppressions(fargateTaskDefinition.executionRole!, [{
        id: 'AwsSolutions-IAM5',
        reason: 'Allow * permissions.',
        appliesTo: ['Resource::*']
    }], true);
}

  2. Import the CDK CI/CD Wrapper module at the top of the file:

    import * as wrapper from '@cdklabs/cdk-cicd-wrapper';
Show Solution

The lib/demo-stack.ts file should look like this: 

import * as cdk from 'aws-cdk-lib';
import { Construct } from 'constructs';
import * as ecs from 'aws-cdk-lib/aws-ecs';
import * as iam from 'aws-cdk-lib/aws-iam';
import * as ec2 from 'aws-cdk-lib/aws-ec2';
import * as ecrAsset from 'aws-cdk-lib/aws-ecr-assets';
import * as ecsPatterns from 'aws-cdk-lib/aws-ecs-patterns';
import * as path from 'path';
import * as nag from 'cdk-nag';
import * as wrapper from '@cdklabs/cdk-cicd-wrapper';

const allowedIPs = ['<your-ip-address>/32'];

export class DemoStack extends cdk.Stack {
  constructor(scope: Construct, id: string, props?: cdk.StackProps) {
    super(scope, id, props);

    // Create an ECS Cluster
    const cluster = new ecs.Cluster(this, 'DemoCluster', {
      clusterName: cdk.Names.uniqueResourceName(this, {
        maxLength: 50,
      }),
      enableFargateCapacityProviders: true,
      containerInsights: true,
    });

    cluster.vpc.addFlowLog('demo-flow-log');
    if (cdk.FeatureFlags.of(this).isEnabled('feature-streamlit')) {
      wrapper.logger.info('Feature Streamlit is enabled');

      const taskRole = new iam.Role(this, 'TaskRole', {
        assumedBy: new iam.ServicePrincipal('ecs-tasks.amazonaws.com'),
        inlinePolicies: {
          BedrockInvokeModel: new iam.PolicyDocument({
            statements: [
              new iam.PolicyStatement({
                actions: ['bedrock:InvokeModel'],
                resources: [`arn:aws:bedrock:${cdk.Aws.REGION}::foundation-model/amazon.titan-image-generator-v1`],
              }),
            ],
          }),
        },
      });

      // Define a Fargate task definition
      const fargateTaskDefinition = new ecs.FargateTaskDefinition(this, 'StreamlitAppTaskDef', {
        taskRole,
      });

      // Add container to task definition
      const container = fargateTaskDefinition.addContainer('StreamlitContainer', {
        image: ecs.ContainerImage.fromAsset(path.resolve(__dirname, 'image_prompts'), { networkMode: ecrAsset.NetworkMode.custom('sagemaker') }),
        memoryLimitMiB: 512,
        cpu: 256,
        logging: new ecs.AwsLogDriver({ streamPrefix: 'streamlit-app' }),
      });

      // Expose the container on port 8501 (Streamlit default port)
      container.addPortMappings({
        containerPort: 8501,
        protocol: ecs.Protocol.TCP,
      });

      // Create security group for ALB
      const albSecurityGroup = new ec2.SecurityGroup(this, 'ALBSecurityGroup', {
        vpc: cluster.vpc,
        allowAllOutbound: true,
        description: 'Security group for ALB with restricted ingress',
      });

      for (const ip of allowedIPs) {
        albSecurityGroup.addIngressRule(
          ec2.Peer.ipv4(ip),
          ec2.Port.tcp(80),
          `Allow HTTP from ${ip}`,
        );
      }

      // Create Fargate service in ECS
      const loadBalancer = new ecsPatterns.ApplicationLoadBalancedFargateService(this, 'StreamlitFargateService', {
        cluster: cluster, // Required
        taskDefinition: fargateTaskDefinition,
        desiredCount: 1,
        publicLoadBalancer: true,
        openListener: false,
      });

      loadBalancer.loadBalancer.addSecurityGroup(albSecurityGroup);

      nag.NagSuppressions.addResourceSuppressions(
        loadBalancer,
        [
          {
            id: 'AwsSolutions-ELB2',
            reason: 'ELB access logs',
          },
        ],
        true,
      );

      nag.NagSuppressions.addResourceSuppressions(fargateTaskDefinition, [{
        id: 'AwsSolutions-IAM5',
        reason: 'Allow * permissions.',
        appliesTo: ['Resource::*'],
      }], true);
    }
  }
}

Great job! The Streamlit app is now controlled by a feature flag.

Step 7: Test the Feature Flag

Let’s test if the feature flag works as expected.

  1. Run the following command to synthesize the CDK code without the feature flag:

    npm run workbench synth

  2. Notice there is no log about the Streamlit app in the console.

  3. Now enable the feature flag by running this command:

    npm run workbench synth -- --context feature-streamlit=true

  4. You should see a log message indicating that the Streamlit app feature is enabled:

    [Info at /sagemaker-user-cdk-cicd-example] Feature Streamlit is enabled

Fantastic! The feature flag works as expected.

Step 8: Enable the Feature Flag by Default in Workbench

If we want to enable the feature flag by default in the Workbench environment, we can set it in the pipeline configuration.

  1. Open the bin/cdk-cicd-example.ts file and update the Workbench configuration:

    .workbench({
    provide(context) {
      context.scope.node.setContext('feature-streamlit', true);
      new DemoStack(context.scope, 'DemoStack', { env: context.environment });
    },
})

Well done! The feature flag is now enabled by default in the Workbench.

Show Solution

The bin/cdk-cicd-example.ts file should look like this: 

#!/usr/bin/env node
import 'source-map-support/register';
import * as cdk from 'aws-cdk-lib';
import * as wrapper from '@cdklabs/cdk-cicd-wrapper';
import { DemoStack } from '../lib/demo-stack';
import { DemoProvider } from './demo-provider';

const app = new cdk.App();

wrapper.PipelineBlueprint.builder()
  .defineStages([
    { stage: wrapper.Stage.RES, account: process.env.AWS_ACCOUNT_ID },
    { stage: wrapper.Stage.DEV, account: process.env.AWS_ACCOUNT_ID },
  ])
  .workbench({
    provide(context) {
      context.scope.node.setContext('feature-streamlit', true);
      new DemoStack(context.scope, 'DemoStack', { env: context.environment });
    },
  })
  .addStack(new DemoProvider())
  .synth(app);

Step 9: Commit and Push the Changes

  1. Run the following command to fix any linting issues:

    npm run lint -- --fix

  2. Run the following commands to update the license and validate the package.json:

    npm run audit:license -- --fix
npm run validate -- --fix

  3. Add the changes to Git:

    git add .

  4. Commit the changes with a meaningful message:

    git commit -m "feat: add feature flag for Streamlit"

  5. Push the changes to the repository:

    git push
✓ Congratulations!
You’ve successfully prepared the Streamlit app for deployment using a feature flag.

Click Next to continue to the next section.

Next