Terraform - How to Refactor This Code Block with Nested for_each?

Question

I have several images already saved to ECR and I want to make these ECR images custom Sagemaker Studio images. The catch is there are also several Sagemaker Studio Domains, so we have to create and attach these Studio images for each domain.

Here is what I have so far:

variable "sagemaker_domains" {
  type    = list(string)
  default = ["my_domain", "your_domain"]
}

custom_docker_images = {
    for domain in var.sagemaker_domains :
    domain => [
      {
        ecr_image_name    = "smstudio-custom"
        studio_image_name = "metaflowkernel"
        display_name      = "Metaflow Kernel"
      },
      {
        ecr_image_name    = "studio-sample-default"
        studio_image_name = "mambakernel"
        display_name      = "Mamba Kernel"
      }
    ]
  }

flat_images = flatten([
    for domain, images in local.custom_docker_images :
    [for image in images : merge({ domain = domain }, image)]
  ])

data "aws_ecr_image" "studio_custom_image" {
  for_each = {
    for image in local.flat_images :
    "${image.domain}-${image.studio_image_name}" => image
  }
  repository_name = each.value.ecr_image_name
  most_recent     = true
}

resource "aws_sagemaker_image" "sagemaker_custom_images" {
  for_each = {
    for image in local.flat_images :
    "${image.domain}-${image.studio_image_name}" => image
  }

  image_name   = "${each.value.studio_image_name}-${each.key}"
  display_name = "${each.value.display_name} - ${each.key}"
  role_arn     = ""
}

resource "aws_sagemaker_image_version" "sagemaker_custom_images" {
  for_each   = aws_sagemaker_image.sagemaker_custom_images
  image_name = each.value.id
  base_image = "${data.aws_caller_identity.current.account_id}.dkr.ecr.${data.aws_region.current.name}.amazonaws.com/${lookup(local.flat_images[each.key], "ecr_image_name")}:${data.aws_ecr_image.studio_custom_image[each.key].image_tags[0]}"
}

This is a bit confusing to me because I essentially first need to loop through the custom_docker_images local variable and create a Sagemaker custom image (and image version). One for each domain. Then I also need to pull in the corresponding most recent ECR image using ecr_image_name.

Error Message:

╷
│ Error: Invalid index
│ 
│   on main.tf line 47, in resource "aws_sagemaker_image_version" "sagemaker_custom_images":
│   47:   base_image = "${data.aws_caller_identity.current.account_id}.dkr.ecr.${data.aws_region.current.name}.amazonaws.com/${lookup(local.flat_images[each.key], "ecr_image_name")}:${data.aws_ecr_image.studio_custom_image[each.key].image_tags[0]}"
│     ├────────────────
│     │ each.key is "amp-mambakernel"
│     │ local.flat_images is tuple with 2 elements
│ 
│ The given key does not identify an element in this collection value: a number is required.
╵
╷
│ Error: Invalid index
│ 
│   on main.tf line 47, in resource "aws_sagemaker_image_version" "sagemaker_custom_images":
│   47:   base_image = "${data.aws_caller_identity.current.account_id}.dkr.ecr.${data.aws_region.current.name}.amazonaws.com/${lookup(local.flat_images[each.key], "ecr_image_name")}:${data.aws_ecr_image.studio_custom_image[each.key].image_tags[0]}"
│     ├────────────────
│     │ each.key is "amp-metaflowkernel"
│     │ local.flat_images is tuple with 2 elements
│ 
│ The given key does not identify an element in this collection value: a number is required.

Thanks for your guidance!

Answer 1

The error you have is due to the fact that your local value flat_images is a list (thus indexed 0, 1, 2, ...) while you are trying to retrieve the list's values with each.key which in your case is a string (domain-studio_image_name).

I would suggest you to turn flat_images into a map, so that you can retrieve its values using the key you choose. For example, you could change flat_images like this:

flat_images = {
    for entry in flatten([
      for domain, images in local.custom_docker_images : [
        for image in images : {
          key   = "${domain}-${image.studio_image_name}"
          value = image
      }]
  ]) : entry.key => entry.value }

For reference, the output I get when I print the value of flat_images obtained as described above, using the input you provided in the question is:

{
  "my_domain-mambakernel" = {
    "display_name" = "Mamba Kernel"
    "ecr_image_name" = "studio-sample-default"
    "studio_image_name" = "mambakernel"
  }
  "my_domain-metaflowkernel" = {
    "display_name" = "Metaflow Kernel"
    "ecr_image_name" = "smstudio-custom"
    "studio_image_name" = "metaflowkernel"
  }
  "your_domain-mambakernel" = {
    "display_name" = "Mamba Kernel"
    "ecr_image_name" = "studio-sample-default"
    "studio_image_name" = "mambakernel"
  }
  "your_domain-metaflowkernel" = {
    "display_name" = "Metaflow Kernel"
    "ecr_image_name" = "smstudio-custom"
    "studio_image_name" = "metaflowkernel"
  }
}

This also allows you to change your for_each with for_each = local.flat_images which is arguably more readable.