Adding support for spot instances

README.md

@@ -96,7 +96,7 @@ You can find a more complete example that uses this module but also includes set
 | Name | Description | Type | Default | Required |
 |------|-------------|------|---------|:--------:|
 | <a name="input_access_list_cidr_blocks"></a> [access\_list\_cidr\_blocks](#input\_access\_list\_cidr\_blocks) | List of CIDRs we want to grant access to our Metaflow Metadata Service. Usually this is our VPN's CIDR blocks. | `list(string)` | `[]` | no |
-| <a name="input_batch_type"></a> [batch\_type](#input\_batch\_type) | AWS Batch Compute Type ('ec2', 'fargate') | `string` | `"ec2"` | no |
+| <a name="input_batch_type"></a> [batch\_type](#input\_batch\_type) | AWS Batch Compute Type ('ec2', 'ec2\_spot', 'fargate', 'fargate\_spot') | `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 |

modules/computation/README.md

@@ -12,7 +12,7 @@ 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_batch_type"></a> [batch\_type](#input\_batch\_type) | AWS Batch Compute Type ('ec2', 'ec2\_spot', 'fargate', 'fargate\_spot') | `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 |

modules/computation/batch.tf

@@ -59,7 +59,9 @@ resource "aws_batch_compute_environment" "this" {
     ]
 
     # Type of instance Amazon EC2 for on-demand. Can use "SPOT" to use unused instances at discount if available
-    type = local.enable_fargate_on_batch ? "FARGATE" : "EC2"
+    type = local.compute_type
+
+    spot_iam_fleet_role = local.is_spot ? aws_iam_role.spotfleet[0].arn : null
 
     tags = !local.enable_fargate_on_batch ? var.standard_tags : null
   }

modules/computation/iam-batch-execution.tf

@@ -182,3 +182,29 @@ resource "aws_iam_role_policy" "grant_ec2_custom_policies" {
   role   = aws_iam_role.batch_execution_role.name
   policy = data.aws_iam_policy_document.ec2_custom_policies.json
 }
+
+data "aws_iam_policy_document" "spotfleet-assume" {
+  statement {
+    effect  = "Allow"
+    actions = ["sts:AssumeRole"]
+
+    principals {
+      type = "Service"
+
+      identifiers = ["spotfleet.amazonaws.com"]
+    }
+  }
+}
+
+resource "aws_iam_policy_attachment" "spotfleet" {
+  count      = local.is_spot ? 1 : 0
+  name       = "${var.resource_prefix}-spotfleet"
+  policy_arn = "arn:aws:iam::aws:policy/service-role/AmazonEC2SpotFleetTaggingRole"
+  roles      = [aws_iam_role.spotfleet[0].name]
+}
+
+resource "aws_iam_role" "spotfleet" {
+  count              = local.is_spot ? 1 : 0
+  name               = "${var.resource_prefix}-spotfleet"
+  assume_role_policy = data.aws_iam_policy_document.spotfleet-assume.json
+}

modules/computation/locals.tf

@@ -19,4 +19,13 @@ locals {
   ecs_instance_role_name = "${var.resource_prefix}ecs-iam-role${var.resource_suffix}"
 
   enable_fargate_on_batch = var.batch_type == "fargate"
+
+  compute_type_map = {
+    "ec2"          = "EC2"
+    "ec2_spot"     = "SPOT"
+    "fargate"      = "FARGATE"
+    "fargate_spot" = "FARGATE_SPOT"
+  }
+  compute_type = local.compute_type_map[var.batch_type]
+  is_spot      = contains(["ec2_spot", "fargate_spot"], var.batch_type)
 }

modules/computation/variables.tf

@@ -1,7 +1,11 @@
 variable "batch_type" {
   type        = string
-  description = "AWS Batch Compute Type ('ec2', 'fargate')"
+  description = "AWS Batch Compute Type ('ec2', 'ec2_spot', 'fargate', 'fargate_spot')"
   default     = "ec2"
+  validation {
+    condition     = contains(["ec2", "ec2_spot", "fargate", "fargate_spot"], var.batch_type)
+    error_message = "Must be one of 'ec2', 'ec2_spot', 'fargate', 'fargate_spot'"
+  }
 }
 
 variable "compute_environment_desired_vcpus" {

variables.tf

@@ -6,8 +6,12 @@ variable "access_list_cidr_blocks" {
 
 variable "batch_type" {
   type        = string
-  description = "AWS Batch Compute Type ('ec2', 'fargate')"
+  description = "AWS Batch Compute Type ('ec2', 'ec2_spot', 'fargate', 'fargate_spot')"
   default     = "ec2"
+  validation {
+    condition     = contains(["ec2", "ec2_spot", "fargate", "fargate_spot"], var.batch_type)
+    error_message = "Must be one of 'ec2', 'ec2_spot', 'fargate', 'fargate_spot'"
+  }
 }
 
 variable "enable_custom_batch_container_registry" {