Setting Up AWS Bedrock with Terraform: A Comprehensive Guide

Setting Up AWS Bedrock with Terraform: A Comprehensive Guide

Introduction

AWS Bedrock is Amazon's fully managed service that makes foundation models from leading AI companies accessible via an API. For organizations looking to integrate generative AI capabilities into their applications, AWS Bedrock provides a streamlined path to leverage powerful AI models. When combined with Terraform, an Infrastructure as Code (IaC) tool, setting up and managing AWS Bedrock resources becomes consistent, repeatable, and version-controlled.

This guide walks you through using Terraform to set up AWS Bedrock with a focus on custom models, which allow you to fine-tune foundation models for your specific use cases.

Prerequisites

  • AWS account with appropriate permissions
  • Terraform installed (version 1.0.0+)
  • Basic knowledge of AWS services and Terraform
  • AWS CLI configured

Understanding AWS Bedrock Custom Models

Before diving into the Terraform configuration, let's understand what a custom model in AWS Bedrock is:

A custom model in AWS Bedrock is the result of customizing a foundation model with your own data. This customization happens in one of two ways:

  • Fine-tuning: Teaching the model specific domain knowledge
  • Continued pre-training: Further training the model on a broader dataset

The customization process involves creating a job that takes your training data, applies it to a foundation model, and outputs a custom model you can then use for your specific use cases.

Setting Up AWS Bedrock with Terraform

Step 1: Configure the AWS Provider

First, set up your Terraform configuration with the AWS provider:

terraform {
  required_providers {
    aws = {
      source  = "hashicorp/aws"
      version = "~> 5.0"
    }
  }
}

provider "aws" {
  region = "us-west-2"  # Choose the region where Bedrock is available
}

Step 2: Create S3 Buckets for Training and Output Data

AWS Bedrock needs S3 buckets to store training data and model output:

resource "aws_s3_bucket" "training" {
  bucket = "my-bedrock-training-data-bucket"
}

resource "aws_s3_bucket" "output" {
  bucket = "my-bedrock-model-output-bucket"
}

Step 3: Create an IAM Role for Bedrock

AWS Bedrock needs permissions to access your S3 buckets:

resource "aws_iam_role" "bedrock_role" {
  name = "bedrock-custom-model-role"

  assume_role_policy = jsonencode({
    Version = "2012-10-17"
    Statement = [
      {
        Action = "sts:AssumeRole"
        Effect = "Allow"
        Principal = {
          Service = "bedrock.amazonaws.com"
        }
      },
    ]
  })
}

resource "aws_iam_role_policy" "bedrock_s3_access" {
  name = "bedrock-s3-access"
  role = aws_iam_role.bedrock_role.id

  policy = jsonencode({
    Version = "2012-10-17"
    Statement = [
      {
        Action = [
          "s3:GetObject",
          "s3:ListBucket",
          "s3:PutObject"
        ]
        Effect = "Allow"
        Resource = [
          aws_s3_bucket.training.arn,
          "${aws_s3_bucket.training.arn}/*",
          aws_s3_bucket.output.arn,
          "${aws_s3_bucket.output.arn}/*"
        ]
      }
    ]
  })
}

Step 4: Reference a Foundation Model

Before creating a custom model, you need to reference a foundation model:

data "aws_bedrock_foundation_model" "base_model" {
  model_id = "amazon.titan-text-express-v1"  # Choose an appropriate model ID
}

Step 5: Create a Custom Model Resource

Now, create the AWS Bedrock custom model resource:

resource "aws_bedrock_custom_model" "my_model" {
  custom_model_name     = "my-custom-model"
  job_name              = "my-custom-model-job"
  base_model_identifier = data.aws_bedrock_foundation_model.base_model.model_arn
  role_arn              = aws_iam_role.bedrock_role.arn
  
  hyperparameters = {
    "epochCount"              = "3"
    "batchSize"               = "1"
    "learningRate"            = "0.005"
    "learningRateWarmupSteps" = "0"
  }
  
  output_data_config {
    s3_uri = "s3://${aws_s3_bucket.output.id}/output/"
  }
  
  training_data_config {
    s3_uri = "s3://${aws_s3_bucket.training.id}/training/train.jsonl"
  }
  
  # Optional: Validation data
  validation_data_config {
    s3_uri = "s3://${aws_s3_bucket.training.id}/validation/validation.jsonl"
  }
  
  # Optional: Customize encryption
  # custom_model_kms_key_id = aws_kms_key.bedrock_key.arn
  
  # Optional: Specify VPC configuration if needed
  # vpc_config {
  #   security_group_ids = [aws_security_group.bedrock_sg.id]
  #   subnet_ids         = [aws_subnet.bedrock_subnet.id]
  # }
  
  tags = {
    Environment = "development"
    Project     = "AI-Initiative"
  }
}

Step 6: Define Outputs

Add outputs to your Terraform configuration to easily access important information:

output "custom_model_arn" {
  description = "The ARN of the custom model"
  value       = aws_bedrock_custom_model.my_model.custom_model_arn
}

output "job_status" {
  description = "The status of the customization job"
  value       = aws_bedrock_custom_model.my_model.job_status
}

output "training_metrics" {
  description = "Metrics from the training process"
  value       = aws_bedrock_custom_model.my_model.training_metrics
}

Preparing the Training Data

AWS Bedrock expects training data in JSONL format. Each line represents a prompt-completion pair. Here's an example:

{"prompt": "What is cloud computing?", "completion": "Cloud computing is the delivery of computing services over the internet..."}
{"prompt": "Explain AWS services", "completion": "AWS offers a wide range of cloud services including compute, storage..."}

Upload this file to your training S3 bucket at the location specified in your Terraform configuration.

Managing the Lifecycle of Your Custom Model

Customization Job Status

After applying your Terraform configuration, the customization job will start running. The job typically takes several hours to complete, depending on the size of your training data and the hyperparameters you've chosen.

You can check the status using:

terraform output job_status

A successful job transitions from InProgress to Completed when your custom model is ready to use.

Updating Your Custom Model

To update your custom model, you'll need to:

  1. Update your training data
  2. Change the job_name in your Terraform configuration
  3. Apply the changes

Note that you cannot update an existing customization job. Instead, Terraform will create a new job and replace the old one.

Deleting Your Custom Model

When you run terraform destroy, Terraform will:

  1. Stop the customization job if it's still active
  2. Delete the custom model if the job has completed

You can customize the timeout for the deletion process:

resource "aws_bedrock_custom_model" "my_model" {
  # ... other configuration ...
  
  timeouts {
    delete = "180m"  # Increase from the default 120m if needed
  }
}

Common Issues and Troubleshooting

  • Access Denied Errors: Ensure your IAM role has the necessary permissions for S3 and Bedrock.
  • Invalid Training Data: Verify your JSONL format is correct and accessible in the S3 location.
  • Failed Jobs: Check the job status and metrics for details on why a job might have failed.
  • Timeout During Creation: Increase the timeout for the resource if your job takes longer than expected.

Advanced Configuration

Using VPC for Enhanced Security

For enhanced security, you can run your customization job within a VPC:

resource "aws_security_group" "bedrock_sg" {
  name        = "bedrock-custom-model-sg"
  description = "Security group for Bedrock custom model"
  vpc_id      = aws_vpc.main.id
  
  egress {
    from_port   = 0
    to_port     = 0
    protocol    = "-1"
    cidr_blocks = ["0.0.0.0/0"]
  }
}

resource "aws_bedrock_custom_model" "my_model" {
  # ... other configuration ...
  
  vpc_config {
    security_group_ids = [aws_security_group.bedrock_sg.id]
    subnet_ids         = [aws_subnet.private_subnet.id]
  }
}

Custom Encryption with KMS

To encrypt your custom model with your own KMS key:

resource "aws_kms_key" "bedrock_key" {
  description             = "KMS key for Bedrock custom model"
  deletion_window_in_days = 7
}

resource "aws_bedrock_custom_model" "my_model" {
  # ... other configuration ...
  
  custom_model_kms_key_id = aws_kms_key.bedrock_key.arn
}

Conclusion

Setting up AWS Bedrock using Terraform provides a robust and repeatable approach to managing your custom AI models. By leveraging the power of Infrastructure as Code, you can version control your model configurations, easily track changes, and integrate model management into your CI/CD pipelines.

As AWS Bedrock continues to evolve, Terraform's AWS provider will likely add more resources and capabilities for managing generative AI infrastructure. Stay up to date with the latest provider documentation to take advantage of new features as they become available.

Additional Resources


For fast and seamless deployment of Terraform in your CI pipeline, Digger is one of the most secure options available.

Check out the repository at
https://github.com/diggerhq/digger to learn more.