Add possibility to use own custom arn roles instead of create new one…

iam.tf

@@ -23,6 +23,8 @@ resource "aws_iam_role" "batch_s3_task_role" {
   assume_role_policy = data.aws_iam_policy_document.batch_s3_task_role_assume_role.json
 
   tags = var.tags
+
+  count = var.batch_s3_task_role_name == "" ? 1 : 0
 }
 
 data "aws_iam_policy_document" "custom_s3_list_batch" {
@@ -202,49 +204,63 @@ data "aws_iam_policy_document" "cloudwatch" {
 
 resource "aws_iam_role_policy" "grant_custom_s3_list_batch" {
   name   = "s3_list"
-  role   = aws_iam_role.batch_s3_task_role.name
+  role   = aws_iam_role.batch_s3_task_role[0].name
   policy = data.aws_iam_policy_document.custom_s3_list_batch.json
+
+  count = var.batch_s3_task_role_name == "" ? 1 : 0
 }
 
 resource "aws_iam_role_policy" "grant_custom_s3_batch" {
   name   = "custom_s3"
-  role   = aws_iam_role.batch_s3_task_role.name
+  role   = aws_iam_role.batch_s3_task_role[0].name
   policy = data.aws_iam_policy_document.custom_s3_batch.json
+
+  count = var.batch_s3_task_role_name == "" ? 1 : 0
 }
 
 resource "aws_iam_role_policy" "grant_s3_kms" {
   name   = "s3_kms"
-  role   = aws_iam_role.batch_s3_task_role.name
+  role   = aws_iam_role.batch_s3_task_role[0].name
   policy = data.aws_iam_policy_document.s3_kms.json
+
+  count = var.batch_s3_task_role_name == "" ? 1 : 0
 }
 
 resource "aws_iam_role_policy" "grant_deny_presigned_batch" {
   name   = "deny_presigned"
-  role   = aws_iam_role.batch_s3_task_role.name
+  role   = aws_iam_role.batch_s3_task_role[0].name
   policy = data.aws_iam_policy_document.deny_presigned_batch.json
+
+  count = var.batch_s3_task_role_name == "" ? 1 : 0
 }
 
 resource "aws_iam_role_policy" "grant_allow_sagemaker" {
   name   = "sagemaker"
-  role   = aws_iam_role.batch_s3_task_role.name
+  role   = aws_iam_role.batch_s3_task_role[0].name
   policy = data.aws_iam_policy_document.allow_sagemaker.json
+
+  count = var.batch_s3_task_role_name == "" ? 1 : 0
 }
 
 resource "aws_iam_role_policy" "grant_iam_pass_role" {
   name   = "iam_pass_role"
-  role   = aws_iam_role.batch_s3_task_role.name
+  role   = aws_iam_role.batch_s3_task_role[0].name
   policy = data.aws_iam_policy_document.iam_pass_role.json
+
+  count = var.batch_s3_task_role_name == "" ? 1 : 0
 }
 
 resource "aws_iam_role_policy" "grant_dynamodb" {
-  count  = var.enable_step_functions ? 1 : 0
+  count  = (var.batch_s3_task_role_name == "" && var.enable_step_functions) ? 1 : 0
   name   = "dynamodb"
-  role   = aws_iam_role.batch_s3_task_role.name
+  role   = aws_iam_role.batch_s3_task_role[0].name
   policy = data.aws_iam_policy_document.dynamodb.json
 }
 
 resource "aws_iam_role_policy" "grant_cloudwatch" {
   name   = "cloudwatch"
-  role   = aws_iam_role.batch_s3_task_role.name
+  role   = aws_iam_role.batch_s3_task_role[0].name
   policy = data.aws_iam_policy_document.cloudwatch.json
+
+  count = var.batch_s3_task_role_name == "" ? 1 : 0
 }

locals.tf

@@ -2,6 +2,11 @@ module "metaflow-common" {
   source = "./modules/common"
 }
 
+data "aws_iam_role" "batch_s3_task_role" {
+  name  = var.batch_s3_task_role_name
+  count = var.batch_s3_task_role_name == "" ? 0 : 1
+}
+
 locals {
   resource_prefix = length(var.resource_prefix) > 0 ? "${var.resource_prefix}-" : ""
   resource_suffix = length(var.resource_suffix) > 0 ? "-${var.resource_suffix}" : ""
@@ -21,4 +26,10 @@ locals {
     module.metaflow-common.default_ui_static_container_image :
     var.ui_static_container_image
   )
+
+  metadata_svc_ecs_task_role_id  = var.metadata_svc_ecs_task_role_name == "" ? aws_iam_role.metadata_svc_ecs_task_role[0].id : data.metadata_svc_ecs_task_role.id
+  metadata_svc_ecs_task_role_arn = var.metadata_svc_ecs_task_role_name == "" ? aws_iam_role.metadata_svc_ecs_task_role[0].arn : data.metadata_svc_ecs_task_role.arn
+
+  batch_s3_task_role_id  = var.batch_s3_task_role_name == "" ? aws_iam_role.batch_s3_task_role[0].id : data.batch_s3_task_role.id
+  batch_s3_task_role_arn = var.batch_s3_task_role_name == "" ? aws_iam_role.batch_s3_task_role[0].arn : data.batch_s3_task_role.arn
 }

main.tf

@@ -43,6 +43,9 @@ module "metaflow-metadata-service" {
   vpc_cidr_blocks                  = var.vpc_cidr_blocks
   with_public_ip                   = var.with_public_ip
 
+  metadata_svc_ecs_task_role_name = var.metadata_svc_ecs_task_role_name
+  lambda_ecs_execute_role_name    = var.lambda_ecs_execute_role_name
+
   standard_tags = var.tags
 }
 
@@ -73,6 +76,8 @@ module "metaflow-ui" {
   certificate_arn                    = var.ui_certificate_arn
   metadata_service_security_group_id = module.metaflow-metadata-service.metadata_service_security_group_id
 
+  metadata_ui_ecs_task_role_name = var.metadata_ui_ecs_task_role_name
+
   extra_ui_static_env_vars  = var.extra_ui_static_env_vars
   extra_ui_backend_env_vars = var.extra_ui_backend_env_vars
   standard_tags             = var.tags
@@ -98,6 +103,10 @@ module "metaflow-computation" {
   launch_template_http_tokens                 = var.launch_template_http_tokens
   launch_template_http_put_response_hop_limit = var.launch_template_http_put_response_hop_limit
 
+  batch_execution_role_name = var.batch_execution_role_name
+  ecs_execution_role_name   = var.ecs_execution_role_name
+  ecs_instance_role_name    = var.ecs_instance_role_name
+
   standard_tags = var.tags
 }
 
@@ -113,5 +122,8 @@ module "metaflow-step-functions" {
   s3_bucket_arn       = module.metaflow-datastore.s3_bucket_arn
   s3_bucket_kms_arn   = module.metaflow-datastore.datastore_s3_bucket_kms_key_arn
 
+  eventbridge_role_name    = var.eventbridge_role_name
+  step_functions_role_name = var.step_functions_role_name
+
   standard_tags = var.tags
 }

modules/computation/batch.tf

@@ -9,7 +9,7 @@ resource "aws_batch_compute_environment" "this" {
   compute_environment_name_prefix = local.compute_env_prefix_name
 
   # Give permissions so the batch service can make API calls.
-  service_role = aws_iam_role.batch_execution_role.arn
+  service_role = local.batch_execution_role_arn
   type         = "MANAGED"
 
   # On destroy, this avoids removing these policies below until compute environments are destroyed
@@ -22,7 +22,7 @@ resource "aws_batch_compute_environment" "this" {
 
   compute_resources {
     # Give permissions so the ECS container instances can make API call.
-    instance_role = !local.enable_fargate_on_batch ? aws_iam_instance_profile.ecs_instance_role.arn : null
+    instance_role = !local.enable_fargate_on_batch ? local.ecs_instance_role_arn : null
 
     # List of types that can be launched.
     instance_type = !local.enable_fargate_on_batch ? var.compute_environment_instance_types : null

modules/computation/ec2.tf

@@ -46,7 +46,9 @@ resource "aws_launch_template" "cpu" {
 */
 resource "aws_iam_instance_profile" "ecs_instance_role" {
   name = local.ecs_instance_role_name
-  role = aws_iam_role.ecs_instance_role.name
+  role = aws_iam_role.ecs_instance_role[0].name
+
+  count = var.ecs_instance_role_name == "" ? 1 : 0  
 }
 
 resource "aws_security_group" "this" {

modules/computation/iam-batch-execution.tf

@@ -24,6 +24,8 @@ resource "aws_iam_role" "batch_execution_role" {
   assume_role_policy = data.aws_iam_policy_document.batch_execution_role_assume_role.json
 
   tags = var.standard_tags
+
+  count = var.batch_execution_role_name == "" ? 1 : 0
 }
 
 data "aws_iam_policy_document" "iam_pass_role" {
@@ -161,24 +163,32 @@ data "aws_iam_policy_document" "ec2_custom_policies" {
 
 resource "aws_iam_role_policy" "grant_iam_pass_role" {
   name   = "iam_pass_role"
-  role   = aws_iam_role.batch_execution_role.name
+  role   = aws_iam_role.batch_execution_role[0].name
   policy = data.aws_iam_policy_document.iam_pass_role.json
+
+  count = var.batch_execution_role_name == "" ? 1 : 0
 }
 
 resource "aws_iam_role_policy" "grant_custom_access_policy" {
   name   = "custom_access"
-  role   = aws_iam_role.batch_execution_role.name
+  role   = aws_iam_role.batch_execution_role[0].name
   policy = data.aws_iam_policy_document.custom_access_policy.json
+
+  count = var.batch_execution_role_name == "" ? 1 : 0
 }
 
 resource "aws_iam_role_policy" "grant_iam_custom_policies" {
   name   = "iam_custom"
-  role   = aws_iam_role.batch_execution_role.name
+  role   = aws_iam_role.batch_execution_role[0].name
   policy = data.aws_iam_policy_document.iam_custom_policies.json
+
+  count = var.batch_execution_role_name == "" ? 1 : 0
 }
 
 resource "aws_iam_role_policy" "grant_ec2_custom_policies" {
   name   = "ec2_custom"
-  role   = aws_iam_role.batch_execution_role.name
+  role   = aws_iam_role.batch_execution_role[0].name
   policy = data.aws_iam_policy_document.ec2_custom_policies.json
+
+  count = var.batch_execution_role_name == "" ? 1 : 0
 }

modules/computation/iam-ecs-execution.tf

@@ -25,6 +25,8 @@ resource "aws_iam_role" "ecs_execution_role" {
   assume_role_policy = data.aws_iam_policy_document.ecs_execution_role_assume_role.json
 
   tags = var.standard_tags
+
+  count = var.ecs_execution_role_name == "" ? 1 : 0  
 }
 
 data "aws_iam_policy_document" "ecs_task_execution_policy" {
@@ -50,6 +52,8 @@ data "aws_iam_policy_document" "ecs_task_execution_policy" {
 
 resource "aws_iam_role_policy" "grant_ecs_access" {
   name   = "ecs_access"
-  role   = aws_iam_role.ecs_execution_role.name
+  role   = aws_iam_role.ecs_execution_role[0].name
   policy = data.aws_iam_policy_document.ecs_task_execution_policy.json
+
+  count = var.ecs_execution_role_name == "" ? 1 : 0  
 }

modules/computation/iam-ecs-instance.tf

@@ -22,6 +22,8 @@ resource "aws_iam_role" "ecs_instance_role" {
   description = "This role is passed to AWS Batch as a `instance_role`. This allows our Metaflow Batch jobs to execute with proper permissions."
 
   assume_role_policy = data.aws_iam_policy_document.ecs_instance_role_assume_role.json
+
+  count = var.ecs_instance_role_name == "" ? 1 : 0  
 }
 
 /*
@@ -32,6 +34,8 @@ resource "aws_iam_role" "ecs_instance_role" {
  https://docs.aws.amazon.com/AmazonECS/latest/developerguide/instance_IAM_role.html
 */
 resource "aws_iam_role_policy_attachment" "ecs_instance_role" {
-  role       = aws_iam_role.ecs_instance_role.name
+  role       = aws_iam_role.ecs_instance_role[0].name
   policy_arn = "arn:${var.iam_partition}:iam::aws:policy/service-role/AmazonEC2ContainerServiceforEC2Role"
+
+  count = var.ecs_instance_role_name == "" ? 1 : 0  
 }

modules/computation/locals.tf

@@ -1,3 +1,13 @@
+data "aws_iam_role" "batch_execution_role" {
+  name = var.batch_execution_role_name
+  count = var.batch_execution_role_name == "" ? 0 : 1
+}
+
+data "aws_iam_role" "ecs_execution_role" {
+  name = var.ecs_execution_role_name
+  count = var.ecs_execution_role_name == "" ? 0 : 1
+}
+
 locals {
   # Name of Batch service's security group used on the compute environment
   batch_security_group_name = "${var.resource_prefix}batch-compute-environment-security-group${var.resource_suffix}"
@@ -19,4 +29,13 @@ locals {
   ecs_instance_role_name = "${var.resource_prefix}ecs-iam-role${var.resource_suffix}"
 
   enable_fargate_on_batch = var.batch_type == "fargate"
+
+  batch_execution_role_id = var.batch_execution_role_name == "" ? aws_iam_role.batch_execution_role[0].id : data.batch_execution_role.id
+  batch_execution_role_arn = var.batch_execution_role_name == "" ? aws_iam_role.batch_execution_role[0].arn : data.batch_execution_role.arn
+
+  ecs_execution_role_id = var.ecs_execution_role_name == "" ? aws_iam_role.ecs_execution_role[0].id : data.ecs_execution_role.id
+  ecs_execution_role_arn = var.ecs_execution_role_name == "" ? aws_iam_role.ecs_execution_role[0].arn : data.ecs_execution_role.arn
+
+  ecs_instance_role_id = var.ecs_instance_role_name == "" ? aws_iam_role.ecs_instance_role[0].id : data.ecs_instance_role.id
+  ecs_instance_role_arn = var.ecs_instance_role_name == "" ? aws_iam_role.ecs_instance_role[0].arn : data.ecs_instance_role.arn
 }

modules/computation/outputs.tf

@@ -9,12 +9,12 @@ output "batch_job_queue_arn" {
 }
 
 output "ecs_execution_role_arn" {
-  value       = aws_iam_role.ecs_execution_role.arn
+  value       = local.ecs_execution_role_arn
   description = "The IAM role that grants access to ECS and Batch services which we'll use as our Metadata Service API's execution_role for our Fargate instance"
 }
 
 output "ecs_instance_role_arn" {
-  value       = aws_iam_role.ecs_instance_role.arn
+  value       = local.ecs_instance_role_arn
   description = "This role will be granted access to our S3 Bucket which acts as our blob storage."
 }
 

modules/computation/variables.tf

@@ -102,3 +102,21 @@ variable "launch_template_image_id" {
   nullable    = true
   default     = null
 }
+
+variable "batch_execution_role_name" {
+  type        = string
+  description = "Custom Name for the Batch Execution Role"
+  default     = ""
+}
+
+variable "ecs_execution_role_name" {
+  type        = string
+  description = "Custom Name for the ECS Execution Role"
+  default     = ""
+}
+
+variable "ecs_instance_role_name" {
+  type        = string
+  description = "Custom Name for the ECS Instance Role"
+  default     = ""
+}

modules/metadata-service/iam.tf

@@ -22,6 +22,8 @@ resource "aws_iam_role" "metadata_svc_ecs_task_role" {
   assume_role_policy = data.aws_iam_policy_document.metadata_svc_ecs_task_assume_role.json
 
   tags = var.standard_tags
+
+  count = var.metadata_svc_ecs_task_role_name == "" ? 1 : 0  
 }
 
 data "aws_iam_policy_document" "s3_kms" {
@@ -84,18 +86,24 @@ data "aws_iam_policy_document" "deny_presigned_batch" {
 
 resource "aws_iam_role_policy" "grant_s3_kms" {
   name   = "s3_kms"
-  role   = aws_iam_role.metadata_svc_ecs_task_role.name
+  role   = aws_iam_role.metadata_svc_ecs_task_role[0].name
   policy = data.aws_iam_policy_document.s3_kms.json
+
+  count = var.metadata_svc_ecs_task_role_name == "" ? 1 : 0  
 }
 
 resource "aws_iam_role_policy" "grant_custom_s3_batch" {
   name   = "custom_s3"
-  role   = aws_iam_role.metadata_svc_ecs_task_role.name
+  role   = aws_iam_role.metadata_svc_ecs_task_role[0].name
   policy = data.aws_iam_policy_document.custom_s3_batch.json
+
+  count = var.metadata_svc_ecs_task_role_name == "" ? 1 : 0 
 }
 
 resource "aws_iam_role_policy" "grant_deny_presigned_batch" {
   name   = "deny_presigned"
-  role   = aws_iam_role.metadata_svc_ecs_task_role.name
+  role   = aws_iam_role.metadata_svc_ecs_task_role[0].name
   policy = data.aws_iam_policy_document.deny_presigned_batch.json
+
+  count = var.metadata_svc_ecs_task_role_name == "" ? 1 : 0 
 }