Support IMDSv2 in launch template (#43)

* Expose variables for IMDSv2

* Enable http_endpoint

* fmt

* pre-commit

* Enable http endpoint by default

* Hop limit defaults to 2

Author: olivermeyer

README.md

@@ -107,6 +107,9 @@ You can find a more complete example that uses this module but also includes set
 | <a name="input_extra_ui_backend_env_vars"></a> [extra\_ui\_backend\_env\_vars](#input\_extra\_ui\_backend\_env\_vars) | Additional environment variables for UI backend container | `map(string)` | `{}` | no |
 | <a name="input_extra_ui_static_env_vars"></a> [extra\_ui\_static\_env\_vars](#input\_extra\_ui\_static\_env\_vars) | Additional environment variables for UI static app | `map(string)` | `{}` | no |
 | <a name="input_iam_partition"></a> [iam\_partition](#input\_iam\_partition) | IAM Partition (Select aws-us-gov for AWS GovCloud, otherwise leave as is) | `string` | `"aws"` | no |
+| <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_metadata_service_container_image"></a> [metadata\_service\_container\_image](#input\_metadata\_service\_container\_image) | Container image for metadata service | `string` | `""` | no |
 | <a name="input_resource_prefix"></a> [resource\_prefix](#input\_resource\_prefix) | string prefix for all resources | `string` | `"metaflow"` | no |
 | <a name="input_resource_suffix"></a> [resource\_suffix](#input\_resource\_suffix) | string suffix for all resources | `string` | `""` | no |

main.tf

@@ -75,16 +75,19 @@ module "metaflow-computation" {
   resource_prefix = local.resource_prefix
   resource_suffix = local.resource_suffix
 
-  batch_type                             = var.batch_type
-  compute_environment_desired_vcpus      = var.compute_environment_desired_vcpus
-  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
-  iam_partition                          = var.iam_partition
-  metaflow_vpc_id                        = var.vpc_id
-  subnet1_id                             = var.subnet1_id
-  subnet2_id                             = var.subnet2_id
+  batch_type                                  = var.batch_type
+  compute_environment_desired_vcpus           = var.compute_environment_desired_vcpus
+  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
+  iam_partition                               = var.iam_partition
+  metaflow_vpc_id                             = var.vpc_id
+  subnet1_id                                  = var.subnet1_id
+  subnet2_id                                  = var.subnet2_id
+  launch_template_http_endpoint               = var.launch_template_http_endpoint
+  launch_template_http_tokens                 = var.launch_template_http_tokens
+  launch_template_http_put_response_hop_limit = var.launch_template_http_put_response_hop_limit
 
   standard_tags = var.tags
 }

modules/computation/README.md

@@ -19,6 +19,9 @@ To read more, see [the Metaflow docs](https://docs.metaflow.org/metaflow-on-aws/
 | <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 |
 | <a name="input_iam_partition"></a> [iam\_partition](#input\_iam\_partition) | IAM Partition (Select aws-us-gov for AWS GovCloud, otherwise leave as is) | `string` | `"aws"` | no |
+| <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_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/ec2.tf

@@ -25,6 +25,12 @@ resource "aws_launch_template" "cpu" {
     }
   }
 
+  metadata_options {
+    http_endpoint               = var.launch_template_http_endpoint
+    http_tokens                 = var.launch_template_http_tokens
+    http_put_response_hop_limit = var.launch_template_http_put_response_hop_limit
+  }
+
   tags = var.standard_tags
 }
 

modules/computation/variables.tf

@@ -65,3 +65,21 @@ variable "subnet2_id" {
   type        = string
   description = "The second private subnet used for redundancy"
 }
+
+variable "launch_template_http_endpoint" {
+  type        = string
+  description = "Whether the metadata service is available. Can be 'enabled' or 'disabled'"
+  default     = "enabled"
+}
+
+variable "launch_template_http_tokens" {
+  type        = string
+  description = "Whether or not the metadata service requires session tokens, also referred to as Instance Metadata Service Version 2 (IMDSv2). Can be 'optional' or 'required'"
+  default     = "optional"
+}
+
+variable "launch_template_http_put_response_hop_limit" {
+  type        = number
+  description = "The desired HTTP PUT response hop limit for instance metadata requests. Can be an integer from 1 to 64"
+  default     = 2
+}

variables.tf

@@ -67,6 +67,24 @@ variable "compute_environment_egress_cidr_blocks" {
   description = "CIDR blocks to which egress is allowed from the Batch Compute environment's security group"
 }
 
+variable "launch_template_http_endpoint" {
+  type        = string
+  description = "Whether the metadata service is available. Can be 'enabled' or 'disabled'"
+  default     = "enabled"
+}
+
+variable "launch_template_http_tokens" {
+  type        = string
+  description = "Whether or not the metadata service requires session tokens, also referred to as Instance Metadata Service Version 2 (IMDSv2). Can be 'optional' or 'required'"
+  default     = "optional"
+}
+
+variable "launch_template_http_put_response_hop_limit" {
+  type        = number
+  description = "The desired HTTP PUT response hop limit for instance metadata requests. Can be an integer from 1 to 64"
+  default     = 2
+}
+
 variable "iam_partition" {
   type        = string
   default     = "aws"