diff --git a/README.md b/README.md index ffe03ad..0355980 100644 --- a/README.md +++ b/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 | |------|-------------|------|---------|:--------:| | [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 | -| [batch\_type](#input\_batch\_type) | AWS Batch Compute Type ('ec2', 'fargate') | `string` | `"ec2"` | no | +| [batch\_type](#input\_batch\_type) | AWS Batch Compute Type ('ec2', 'ec2\_spot', 'fargate', 'fargate\_spot') | `string` | `"ec2"` | no | | [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 | | [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)` |
[| no | | [compute\_environment\_instance\_types](#input\_compute\_environment\_instance\_types) | The instance types for the compute environment | `list(string)` |
"0.0.0.0/0"
]
[| no | diff --git a/modules/computation/README.md b/modules/computation/README.md index 8cd5712..8635c1d 100644 --- a/modules/computation/README.md +++ b/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 | |------|-------------|------|---------|:--------:| -| [batch\_type](#input\_batch\_type) | AWS Batch Compute Type ('ec2', 'fargate') | `string` | `"ec2"` | no | +| [batch\_type](#input\_batch\_type) | AWS Batch Compute Type ('ec2', 'ec2\_spot', 'fargate', 'fargate\_spot') | `string` | `"ec2"` | no | | [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 | | [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)` |
"c4.large",
"c4.xlarge",
"c4.2xlarge",
"c4.4xlarge",
"c4.8xlarge"
]
[| no | | [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 | diff --git a/modules/computation/batch.tf b/modules/computation/batch.tf index bd3f9af..4de4fa0 100644 --- a/modules/computation/batch.tf +++ b/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 } diff --git a/modules/computation/iam-batch-execution.tf b/modules/computation/iam-batch-execution.tf index 733cab6..9701d40 100644 --- a/modules/computation/iam-batch-execution.tf +++ b/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 +} diff --git a/modules/computation/locals.tf b/modules/computation/locals.tf index bea84cb..db97a21 100644 --- a/modules/computation/locals.tf +++ b/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) } diff --git a/modules/computation/variables.tf b/modules/computation/variables.tf index 7604e2d..d94a108 100644 --- a/modules/computation/variables.tf +++ b/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" { diff --git a/variables.tf b/variables.tf index 9fbc456..a0960d9 100644 --- a/variables.tf +++ b/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" {
"0.0.0.0/0"
]