From 764c33ae736ee3714f97d5a919961d11f0344d38 Mon Sep 17 00:00:00 2001
From: Emterry <123941245+Emterry@users.noreply.github.com>
Date: Wed, 15 Jan 2025 17:22:56 +0000
Subject: [PATCH] Bedrock batch inference (#6524)

* update bedrock permissions

* add role for batch inference

* role and policies for batch inference in bedrock

* remove not needed code

* linter

* linter

* refactor

* linter

* lint
---
 .../tooling-integration-iam-policies.tf       | 70 +++++++++++++++++++
 1 file changed, 70 insertions(+)

diff --git a/terraform/aws/analytical-platform-development/cluster/tooling-integration-iam-policies.tf b/terraform/aws/analytical-platform-development/cluster/tooling-integration-iam-policies.tf
index ed5998de37..6357cdc7e7 100644
--- a/terraform/aws/analytical-platform-development/cluster/tooling-integration-iam-policies.tf
+++ b/terraform/aws/analytical-platform-development/cluster/tooling-integration-iam-policies.tf
@@ -131,3 +131,73 @@ resource "aws_iam_policy" "textract_integration" {
   description = "Permissions needed to allow access to Textract from tooling."
   policy      = data.aws_iam_policy_document.textract_integration.json
 }
+
+##################################################
+# Bedrock Batch Inference
+##################################################
+
+data "aws_iam_policy_document" "bedrock_batch_inference" {
+  statement {
+    sid     = "AllowBedrockAssumeRoleForBatchInference"
+    actions = ["sts:AssumeRole"]
+    condition {
+      test     = "StringEquals"
+      variable = "aws:SourceAccount"
+      values   = [data.aws_caller_identity.current.account_id]
+    }
+    condition {
+      test     = "ArnLike"
+      variable = "aws:SourceArn"
+      values   = ["arn:aws:bedrock:*:${data.aws_caller_identity.current.account_id}:model-invocation-job/*"]
+    }
+    effect = "Allow"
+    principals {
+      type        = "Service"
+      identifiers = ["bedrock.amazonaws.com"]
+    }
+  }
+}
+
+resource "aws_iam_role" "bedrock_batch_inference" {
+  name               = "bedrock-batch-inference-role"
+  description        = "IAM role for AWS Bedrock to perform batch inference tasks as part of model invocation workflows."
+  assume_role_policy = data.aws_iam_policy_document.bedrock_batch_inference.json
+}
+
+resource "aws_iam_role_policy_attachment" "bedrock_batch_inference" {
+  role       = aws_iam_role.bedrock_batch_inference.name
+  policy_arn = aws_iam_policy.bedrock_integration.arn
+}
+
+# Bedrock Batch Inference s3 access
+data "aws_iam_policy_document" "bedrock_batch_inference_s3_access" {
+  statement {
+    sid    = "bedrock-batch-inference-s3-access"
+    effect = "Allow"
+
+    actions = [
+      "s3:GetObject",
+      "s3:PutObject",
+      "s3:ListBucket",
+    ]
+
+    resources = [
+      "arn:aws:s3:::*"
+    ]
+
+    condition {
+      test     = "StringEquals"
+      variable = "aws:ResourceAccount"
+      values = [
+        data.aws_caller_identity.current.account_id
+      ]
+    }
+  }
+}
+
+resource "aws_iam_policy" "bedrock_batch_inference_s3_access" {
+  name        = "bedrock-batch-inference-s3-access"
+  description = "S3 access policy for Bedrock batch inference."
+  policy      = data.aws_iam_policy_document.bedrock_batch_inference_s3_access.json
+}
+