Create dedicated security group for the Compute environment (#24)

* Create dedicated security group for the Compute environment

* Allow egress

* Expose egress CIDR blocks

* Delete vpc_default data source; allow internal traffic

Author: olivermeyer

README.md

@@ -98,6 +98,7 @@ You can find a more complete example that uses this module but also includes set
 | <a name="input_api_basic_auth"></a> [api\_basic\_auth](#input\_api\_basic\_auth) | Enable basic auth for API Gateway? (requires key export) | `bool` | `true` | no |
 | <a name="input_batch_type"></a> [batch\_type](#input\_batch\_type) | AWS Batch Compute Type ('ec2', 'fargate') | `string` | `"ec2"` | no |
 | <a name="input_compute_environment_desired_vcpus"></a> [compute\_environment\_desired\_vcpus](#input\_compute\_environment\_desired\_vcpus) | Desired Starting VCPUs for Batch Compute Environment [0-16] for EC2 Batch Compute Environment (ignored for Fargate) | `number` | `8` | no |
+| <a name="input_compute_environment_egress_cidr_blocks"></a> [compute\_environment\_egress\_cidr\_blocks](#input\_compute\_environment\_egress\_cidr\_blocks) | CIDR blocks to which egress is allowed from the Batch Compute environment's security group | `list(string)` | <pre>[<br>  "0.0.0.0/0"<br>]</pre> | no |
 | <a name="input_compute_environment_instance_types"></a> [compute\_environment\_instance\_types](#input\_compute\_environment\_instance\_types) | The instance types for the compute environment | `list(string)` | <pre>[<br>  "c4.large",<br>  "c4.xlarge",<br>  "c4.2xlarge",<br>  "c4.4xlarge",<br>  "c4.8xlarge"<br>]</pre> | no |
 | <a name="input_compute_environment_max_vcpus"></a> [compute\_environment\_max\_vcpus](#input\_compute\_environment\_max\_vcpus) | Maximum VCPUs for Batch Compute Environment [16-96] | `number` | `64` | no |
 | <a name="input_compute_environment_min_vcpus"></a> [compute\_environment\_min\_vcpus](#input\_compute\_environment\_min\_vcpus) | Minimum VCPUs for Batch Compute Environment [0-16] for EC2 Batch Compute Environment (ignored for Fargate) | `number` | `8` | no |
@@ -133,6 +134,7 @@ You can find a more complete example that uses this module but also includes set
 | <a name="output_METAFLOW_SFN_DYNAMO_DB_TABLE"></a> [METAFLOW\_SFN\_DYNAMO\_DB\_TABLE](#output\_METAFLOW\_SFN\_DYNAMO\_DB\_TABLE) | AWS DynamoDB table name for tracking AWS Step Functions execution metadata. |
 | <a name="output_METAFLOW_SFN_IAM_ROLE"></a> [METAFLOW\_SFN\_IAM\_ROLE](#output\_METAFLOW\_SFN\_IAM\_ROLE) | IAM role for AWS Step Functions to access AWS resources (AWS Batch, AWS DynamoDB). |
 | <a name="output_api_gateway_rest_api_id_key_id"></a> [api\_gateway\_rest\_api\_id\_key\_id](#output\_api\_gateway\_rest\_api\_id\_key\_id) | API Gateway Key ID for Metadata Service. Fetch Key from AWS Console [METAFLOW\_SERVICE\_AUTH\_KEY] |
+| <a name="output_batch_compute_environment_security_group_id"></a> [batch\_compute\_environment\_security\_group\_id](#output\_batch\_compute\_environment\_security\_group\_id) | The ID of the security group attached to the Batch Compute environment. |
 | <a name="output_datastore_s3_bucket_kms_key_arn"></a> [datastore\_s3\_bucket\_kms\_key\_arn](#output\_datastore\_s3\_bucket\_kms\_key\_arn) | The ARN of the KMS key used to encrypt the Metaflow datastore S3 bucket |
 | <a name="output_metadata_svc_ecs_task_role_arn"></a> [metadata\_svc\_ecs\_task\_role\_arn](#output\_metadata\_svc\_ecs\_task\_role\_arn) | n/a |
 | <a name="output_metaflow_api_gateway_rest_api_id"></a> [metaflow\_api\_gateway\_rest\_api\_id](#output\_metaflow\_api\_gateway\_rest\_api\_id) | The ID of the API Gateway REST API we'll use to accept MetaData service requests to forward to the Fargate API instance |

main.tf

@@ -78,6 +78,7 @@ module "metaflow-computation" {
   compute_environment_instance_types      = var.compute_environment_instance_types
   compute_environment_max_vcpus           = var.compute_environment_max_vcpus
   compute_environment_min_vcpus           = var.compute_environment_min_vcpus
+  compute_environment_egress_cidr_blocks  = var.compute_environment_egress_cidr_blocks
   enable_step_functions                   = var.enable_step_functions
   iam_partition                           = var.iam_partition
   metaflow_step_functions_dynamodb_policy = module.metaflow-step-functions.metaflow_step_functions_dynamodb_policy

modules/computation/README.md

@@ -14,6 +14,7 @@ To read more, see [the Metaflow docs](https://docs.metaflow.org/metaflow-on-aws/
 |------|-------------|------|---------|:--------:|
 | <a name="input_batch_type"></a> [batch\_type](#input\_batch\_type) | AWS Batch Compute Type ('ec2', 'fargate') | `string` | `"ec2"` | no |
 | <a name="input_compute_environment_desired_vcpus"></a> [compute\_environment\_desired\_vcpus](#input\_compute\_environment\_desired\_vcpus) | Desired Starting VCPUs for Batch Compute Environment [0-16] for EC2 Batch Compute Environment (ignored for Fargate) | `number` | n/a | yes |
+| <a name="input_compute_environment_egress_cidr_blocks"></a> [compute\_environment\_egress\_cidr\_blocks](#input\_compute\_environment\_egress\_cidr\_blocks) | CIDR blocks to which egress is allowed from the Batch Compute environment's security group | `list(string)` | <pre>[<br>  "0.0.0.0/0"<br>]</pre> | no |
 | <a name="input_compute_environment_instance_types"></a> [compute\_environment\_instance\_types](#input\_compute\_environment\_instance\_types) | The instance types for the compute environment as a comma-separated list | `list(string)` | n/a | yes |
 | <a name="input_compute_environment_max_vcpus"></a> [compute\_environment\_max\_vcpus](#input\_compute\_environment\_max\_vcpus) | Maximum VCPUs for Batch Compute Environment [16-96] | `number` | n/a | yes |
 | <a name="input_compute_environment_min_vcpus"></a> [compute\_environment\_min\_vcpus](#input\_compute\_environment\_min\_vcpus) | Minimum VCPUs for Batch Compute Environment [0-16] for EC2 Batch Compute Environment (ignored for Fargate) | `number` | n/a | yes |
@@ -32,6 +33,7 @@ To read more, see [the Metaflow docs](https://docs.metaflow.org/metaflow-on-aws/
 | Name | Description |
 |------|-------------|
 | <a name="output_METAFLOW_BATCH_JOB_QUEUE"></a> [METAFLOW\_BATCH\_JOB\_QUEUE](#output\_METAFLOW\_BATCH\_JOB\_QUEUE) | AWS Batch Job Queue ARN for Metaflow |
+| <a name="output_batch_compute_environment_security_group_id"></a> [batch\_compute\_environment\_security\_group\_id](#output\_batch\_compute\_environment\_security\_group\_id) | The ID of the security group attached to the Batch Compute environment. |
 | <a name="output_batch_job_queue_arn"></a> [batch\_job\_queue\_arn](#output\_batch\_job\_queue\_arn) | The ARN of the job queue we'll use to accept Metaflow tasks |
 | <a name="output_ecs_execution_role_arn"></a> [ecs\_execution\_role\_arn](#output\_ecs\_execution\_role\_arn) | 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 |
 | <a name="output_ecs_instance_role_arn"></a> [ecs\_instance\_role\_arn](#output\_ecs\_instance\_role\_arn) | This role will be granted access to our S3 Bucket which acts as our blob storage. |

modules/computation/batch.tf

@@ -49,7 +49,7 @@ resource "aws_batch_compute_environment" "this" {
 
     # Security group to apply to the instances launched.
     security_group_ids = [
-      data.aws_security_group.vpc_default.id,
+      aws_security_group.this.id,
     ]
 
     # Which subnet to launch the instances into.

modules/computation/data.tf

@@ -7,8 +7,3 @@ data "aws_ssm_parameter" "ecs_optimized_cpu_ami" {
 data "aws_ssm_parameter" "ecs_optimized_gpu_ami" {
   name = "/aws/service/ecs/optimized-ami/amazon-linux-2/gpu/recommended"
 }
-
-data "aws_security_group" "vpc_default" {
-  name   = "default"
-  vpc_id = var.metaflow_vpc_id
-}

modules/computation/ec2.tf

@@ -41,3 +41,23 @@ resource "aws_iam_instance_profile" "ecs_instance_role" {
   name = local.ecs_instance_role_name
   role = aws_iam_role.ecs_instance_role.name
 }
+
+resource "aws_security_group" "this" {
+  name   = local.batch_security_group_name
+  vpc_id = var.metaflow_vpc_id
+
+  egress {
+    from_port   = 0
+    to_port     = 0
+    protocol    = "-1"
+    cidr_blocks = var.compute_environment_egress_cidr_blocks
+  }
+
+  ingress {
+    from_port   = 0
+    to_port     = 0
+    protocol    = "-1"
+    self        = true
+    description = "internal traffic"
+  }
+}

modules/computation/outputs.tf

@@ -17,3 +17,8 @@ output "ecs_instance_role_arn" {
   value       = aws_iam_role.ecs_instance_role.arn
   description = "This role will be granted access to our S3 Bucket which acts as our blob storage."
 }
+
+output "batch_compute_environment_security_group_id" {
+  value       = aws_security_group.this.id
+  description = "The ID of the security group attached to the Batch Compute environment."
+}

modules/computation/variables.tf

@@ -24,6 +24,12 @@ variable "compute_environment_min_vcpus" {
   description = "Minimum VCPUs for Batch Compute Environment [0-16] for EC2 Batch Compute Environment (ignored for Fargate)"
 }
 
+variable "compute_environment_egress_cidr_blocks" {
+  type        = list(string)
+  default     = ["0.0.0.0/0"]
+  description = "CIDR blocks to which egress is allowed from the Batch Compute environment's security group"
+}
+
 variable "enable_step_functions" {
   default     = false
   description = "If true, apply policies required for step functions"

outputs.tf

@@ -123,3 +123,8 @@ output "ui_alb_arn" {
   value       = (length(module.metaflow-ui) > 0) ? module.metaflow-ui[0].alb_arn : ""
   description = "UI ALB ARN"
 }
+
+output "batch_compute_environment_security_group_id" {
+  value       = module.metaflow-computation.batch_compute_environment_security_group_id
+  description = "The ID of the security group attached to the Batch Compute environment."
+}

variables.tf

@@ -61,6 +61,12 @@ variable "compute_environment_max_vcpus" {
   default     = 64
 }
 
+variable "compute_environment_egress_cidr_blocks" {
+  type        = list(string)
+  default     = ["0.0.0.0/0"]
+  description = "CIDR blocks to which egress is allowed from the Batch Compute environment's security group"
+}
+
 variable "iam_partition" {
   type        = string
   default     = "aws"