Enable some variables to be externally provided to the computation module

modules/computation/README.md

@@ -13,6 +13,8 @@ To read more, see [the Metaflow docs](https://docs.metaflow.org/metaflow-on-aws/
 | Name | Description | Type | Default | Required |
 |------|-------------|------|---------|:--------:|
 | <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_additional_security_group_ids"></a> [compute\_environment\_additional\_security\_group\_ids](#input\_compute\_environment\_additional\_security\_group\_ids) | Additional security group ids to apply to the Batch Compute environment | `list(string)` | `[]` | no |
+| <a name="input_compute_environment_allocation_strategy"></a> [compute\_environment\_allocation\_strategy](#input\_compute\_environment\_allocation\_strategy) | Allocation strategy for Batch Compute environment (BEST\_FIT, BEST\_FIT\_PROGRESSIVE, SPOT\_CAPACITY\_OPTIMIZED) | `string` | `"BEST_FIT"` | 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 |
@@ -22,6 +24,7 @@ To read more, see [the Metaflow docs](https://docs.metaflow.org/metaflow-on-aws/
 | <a name="input_launch_template_http_endpoint"></a> [launch\_template\_http\_endpoint](#input\_launch\_template\_http\_endpoint) | Whether the metadata service is available. Can be 'enabled' or 'disabled' | `string` | `"enabled"` | no |
 | <a name="input_launch_template_http_put_response_hop_limit"></a> [launch\_template\_http\_put\_response\_hop\_limit](#input\_launch\_template\_http\_put\_response\_hop\_limit) | The desired HTTP PUT response hop limit for instance metadata requests. Can be an integer from 1 to 64 | `number` | `2` | no |
 | <a name="input_launch_template_http_tokens"></a> [launch\_template\_http\_tokens](#input\_launch\_template\_http\_tokens) | Whether or not the metadata service requires session tokens, also referred to as Instance Metadata Service Version 2 (IMDSv2). Can be 'optional' or 'required' | `string` | `"optional"` | no |
+| <a name="input_launch_template_image_id"></a> [launch\_template\_image\_id](#input\_launch\_template\_image\_id) | AMI id for launch template, defaults to allow AWS Batch to decide | `string` | `null` | no |
 | <a name="input_metaflow_vpc_id"></a> [metaflow\_vpc\_id](#input\_metaflow\_vpc\_id) | ID of the Metaflow VPC this SageMaker notebook instance is to be deployed in | `string` | n/a | yes |
 | <a name="input_resource_prefix"></a> [resource\_prefix](#input\_resource\_prefix) | Prefix given to all AWS resources to differentiate between applications | `string` | n/a | yes |
 | <a name="input_resource_suffix"></a> [resource\_suffix](#input\_resource\_suffix) | Suffix given to all AWS resources to differentiate between environment and workspace | `string` | n/a | yes |

modules/computation/batch.tf

@@ -33,7 +33,7 @@ resource "aws_batch_compute_environment" "this" {
     desired_vcpus = !local.enable_fargate_on_batch ? var.compute_environment_desired_vcpus : null
 
     # Prefers cheap vCPU approaches
-    allocation_strategy = !local.enable_fargate_on_batch ? "BEST_FIT" : null
+    allocation_strategy = !local.enable_fargate_on_batch ? var.compute_environment_allocation_strategy : null
 
     /* Links to a launch template who has more than the standard 8GB of disk space. So we can download training data.
        Always uses the "default version", which means we can update the Launch Template to a smaller or larger disk size
@@ -48,9 +48,9 @@ resource "aws_batch_compute_environment" "this" {
     }
 
     # Security group to apply to the instances launched.
-    security_group_ids = [
+    security_group_ids = concat([
       aws_security_group.this.id,
-    ]
+    ], var.compute_environment_additional_security_group_ids)
 
     # Which subnet to launch the instances into.
     subnets = [

modules/computation/ec2.tf

@@ -15,6 +15,9 @@ resource "aws_launch_template" "cpu" {
     arn = aws_iam_instance_profile.ecs_instance_role.arn
   }
 
+  # Null image_id allows AWS Batch to decide.
+  image_id = var.launch_template_image_id
+
   block_device_mappings {
     device_name = "/dev/xvda"
 

modules/computation/variables.tf

@@ -30,6 +30,18 @@ variable "compute_environment_egress_cidr_blocks" {
   description = "CIDR blocks to which egress is allowed from the Batch Compute environment's security group"
 }
 
+variable "compute_environment_additional_security_group_ids" {
+  type        = list(string)
+  default     = []
+  description = "Additional security group ids to apply to the Batch Compute environment"
+}
+
+variable "compute_environment_allocation_strategy" {
+  type        = string
+  default     = "BEST_FIT"
+  description = "Allocation strategy for Batch Compute environment (BEST_FIT, BEST_FIT_PROGRESSIVE, SPOT_CAPACITY_OPTIMIZED)"
+}
+
 variable "iam_partition" {
   type        = string
   default     = "aws"
@@ -83,3 +95,10 @@ variable "launch_template_http_put_response_hop_limit" {
   description = "The desired HTTP PUT response hop limit for instance metadata requests. Can be an integer from 1 to 64"
   default     = 2
 }
+
+variable "launch_template_image_id" {
+  type        = string
+  description = "AMI id for launch template, defaults to allow AWS Batch to decide"
+  nullable    = true
+  default     = null
+}