From cae70553dd5fc91bf9b0f0b06325e44e4ecceb41 Mon Sep 17 00:00:00 2001
From: Emterry <123941245+Emterry@users.noreply.github.com>
Date: Fri, 17 Jan 2025 14:14:26 +0000
Subject: [PATCH] adding cross-inference into dev and batch inference into prod
 (#6571)

* adding cross-inference into dev and batch inference into prod

* linter

* add data block

* linter
---
 .../tooling-iam/data.tf                       |  1 +
 .../tooling-integration-iam-policies.tf       | 81 ++++++++++++++++++-
 .../tooling-integration-iam-policies.tf       | 14 ++++
 3 files changed, 95 insertions(+), 1 deletion(-)
 create mode 100644 terraform/aws/analytical-platform-data-production/tooling-iam/data.tf

diff --git a/terraform/aws/analytical-platform-data-production/tooling-iam/data.tf b/terraform/aws/analytical-platform-data-production/tooling-iam/data.tf
new file mode 100644
index 0000000000..8fc4b38cc5
--- /dev/null
+++ b/terraform/aws/analytical-platform-data-production/tooling-iam/data.tf
@@ -0,0 +1 @@
+data "aws_caller_identity" "current" {}
diff --git a/terraform/aws/analytical-platform-data-production/tooling-iam/tooling-integration-iam-policies.tf b/terraform/aws/analytical-platform-data-production/tooling-iam/tooling-integration-iam-policies.tf
index 4b763cde91..455bf65e66 100644
--- a/terraform/aws/analytical-platform-data-production/tooling-iam/tooling-integration-iam-policies.tf
+++ b/terraform/aws/analytical-platform-data-production/tooling-iam/tooling-integration-iam-policies.tf
@@ -49,7 +49,11 @@ data "aws_iam_policy_document" "bedrock_integration" {
       "bedrock:CreateFoundationModelAgreement",
       "bedrock:DeleteFoundationModelAgreement",
       "bedrock:ListFoundationModelAgreementOffers",
-      "bedrock:GetUseCaseForModelAccess"
+      "bedrock:GetUseCaseForModelAccess",
+      "bedrock:CreateModelInvocationJob",
+      "bedrock:GetModelInvocationJob",
+      "bedrock:ListModelInvocationJobs",
+      "bedrock:StopModelInvocationJob"
     ]
     resources = ["*"]
     condition {
@@ -72,6 +76,81 @@ resource "aws_iam_policy" "bedrock_integration" {
   policy      = data.aws_iam_policy_document.bedrock_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    = "BedrockBatchInferenceS3Access"
+    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
+}
+
+resource "aws_iam_role_policy_attachment" "bedrock_batch_inference_s3_access" {
+  role       = aws_iam_role.bedrock_batch_inference.name
+  policy_arn = aws_iam_policy.bedrock_batch_inference_s3_access.arn
+}
+
+
 #tfsec:ignore:aws-iam-no-policy-wildcards
 data "aws_iam_policy_document" "textract_integration" {
   #checkov:skip=CKV_AWS_111: This is a service policy
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 5eb35fbfc7..8b674c989f 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
@@ -54,6 +54,7 @@ data "aws_iam_policy_document" "bedrock_integration" {
       "bedrock:CreateModelInvocationJob",
       "bedrock:GetModelInvocationJob",
       "bedrock:ListModelInvocationJobs",
+      "bedrock:GetInferenceProfiles",
       "bedrock:StopModelInvocationJob"
     ]
 
@@ -156,6 +157,19 @@ data "aws_iam_policy_document" "bedrock_batch_inference" {
       identifiers = ["bedrock.amazonaws.com"]
     }
   }
+  statement {
+    sid    = "CrossRegionInference"
+    effect = "Allow"
+
+    actions = [
+      "bedrock:InvokeModel"
+    ]
+
+    resources = [
+      "arn:aws:bedrock:*::inference-profile/*",
+      "arn:aws:bedrock:*::foundation-model/*"
+    ]
+  }
 }
 
 resource "aws_iam_role" "bedrock_batch_inference" {