feat: Add pattern for Neuron using Trainium and new AL2023 Neuron AMI (…

…#2034)

docs/patterns/machine-learning/aws-neuron-efa.md

@@ -0,0 +1,7 @@
+---
+title: AWS Neuron with EFA
+---
+
+{%
+   include-markdown "../../../patterns/aws-neuron-efa/README.md"
+%}

patterns/aws-neuron-efa/README.md

@@ -0,0 +1,54 @@
+# EKS Cluster w/ AWS Neuron Devices and EFA for Machine Learning
+
+This pattern demonstrates an Amazon EKS Cluster with an EFA-enabled nodegroup that utilizes `trn1.32xlarge` instances that are used in distributed, multi-node machine learning workloads.
+
+The following components are demonstrated in this pattern:
+
+- A "default" node group that supports addons and components that do not require AWS Neuron nor EFA devices. Any pods that do not tolerate the taints of the Neuron node group will be scheduled on instances within this node group.
+- A node group of `trn1.32xlarge` instances with:
+    - all x8 [EFA network interfaces](https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/efa.html) enabled
+    - provisioned within a placement group so that the instances are co-located close to one another in a single availability zone that supports the instance type
+    - a common taint of `"aws.amazon.com/neuron:NoSchedule"` to ensure only the intended applications are permitted to run on the nodes created
+    - two labels identifying that this nodegroup supports AWS Neuron and EFA devices; allowing pods to use node selectors with these labels
+    - the NVME instance store volumes are mounted in a RAID-0 array to provide a single, large, high-performance storage volume for the Neuron workloads
+    - kubelet and containerd are configured to utilize the RAID-0 volume, allowing kubelet to discover the additional storage as ephemeral storage that can be utilized by pods
+- A Helm chart deployment for the [Neuron device plugin](https://github.com/aws-neuron/neuron-helm-charts/tree/main/charts/neuron-helm-chart) to expose and mount the Neuron devices provided by the instances to the pods that request them
+- A Helm chart deployment for the EFA device plugin to expose and mount the EFA network interfaces provided by the instances to the pods that request them. Since the EFA network interfaces are only found on the instances that provide AWS Neuron devices in this pattern, we do not apply an additional taint for the EFA network interfaces to avoid over-constraining.
+
+## Code
+
+```terraform hl_lines="26-28 34-80"
+{% include  "../../patterns/aws-neuron-efa/eks.tf" %}
+```
+
+```terraform hl_lines="9-50"
+{% include  "../../patterns/aws-neuron-efa/helm.tf" %}
+```
+
+## Deploy
+
+See [here](https://aws-ia.github.io/terraform-aws-eks-blueprints/getting-started/#prerequisites) for the prerequisites and steps to deploy this pattern.
+
+## Validate
+
+1. List the nodes and their instance type:
+
+    ```sh
+    kubectl get nodes -L node.kubernetes.io/instance-type
+    ```
+
+    ```text
+    NAME                                        STATUS   ROLES    AGE   VERSION               INSTANCE-TYPE
+    ip-10-0-12-200.us-east-2.compute.internal   Ready    <none>   82m   v1.31.0-eks-a737599   m5.large
+    ip-10-0-24-248.us-east-2.compute.internal   Ready    <none>   82m   v1.31.0-eks-a737599   m5.large
+    ip-10-0-39-213.us-east-2.compute.internal   Ready    <none>   75m   v1.31.0-eks-a737599   trn1.32xlarge
+    ip-10-0-43-172.us-east-2.compute.internal   Ready    <none>   75m   v1.31.0-eks-a737599   trn1.32xlarge
+    ```
+
+    You should see two EFA-enabled (in this example `trn1.32xlarge`) nodes in the list.
+
+## Destroy
+
+{%
+   include-markdown "../../docs/_partials/destroy.md"
+%}

patterns/aws-neuron-efa/eks.tf

@@ -0,0 +1,93 @@
+################################################################################
+# Cluster
+################################################################################
+
+module "eks" {
+  source  = "terraform-aws-modules/eks/aws"
+  version = "~> 20.26"
+
+  cluster_name    = local.name
+  cluster_version = "1.31"
+
+  # Give the Terraform identity admin access to the cluster
+  # which will allow it to deploy resources into the cluster
+  enable_cluster_creator_admin_permissions = true
+  cluster_endpoint_public_access           = true
+
+  cluster_addons = {
+    coredns                = {}
+    eks-pod-identity-agent = {}
+    kube-proxy             = {}
+    vpc-cni = {
+      most_recent = true
+    }
+  }
+
+  # Add security group rules on the node group security group to
+  # allow EFA traffic
+  enable_efa_support = true
+
+  vpc_id     = module.vpc.vpc_id
+  subnet_ids = module.vpc.private_subnets
+
+  eks_managed_node_groups = {
+    neuron-efa = {
+      # The EKS AL2023 Neuron AMI provides all of the necessary components
+      # for accelerated workloads w/ EFA
+      ami_type       = "AL2023_x86_64_NEURON"
+      instance_types = ["trn1.32xlarge"]
+
+      # Mount instance store volumes in RAID-0 for kubelet and containerd
+      # https://github.com/awslabs/amazon-eks-ami/blob/master/doc/USER_GUIDE.md#raid-0-for-kubelet-and-containerd-raid0
+      cloudinit_pre_nodeadm = [
+        {
+          content_type = "application/node.eks.aws"
+          content      = <<-EOT
+            ---
+            apiVersion: node.eks.aws/v1alpha1
+            kind: NodeConfig
+            spec:
+              instance:
+                localStorage:
+                  strategy: RAID0
+          EOT
+        }
+      ]
+
+      min_size     = 2
+      max_size     = 2
+      desired_size = 2
+
+      # This will:
+      # 1. Create a placement group to place the instances close to one another
+      # 2. Ignore subnets that reside in AZs that do not support the instance type
+      # 3. Expose all of the available EFA interfaces on the launch template
+      enable_efa_support = true
+
+      labels = {
+        "vpc.amazonaws.com/efa.present" = "true"
+        "aws.amazon.com/neuron.present" = "true"
+      }
+
+      taints = {
+        # Ensure only Neuron workloads are scheduled on this node group
+        gpu = {
+          key    = "aws.amazon.com/neuron"
+          value  = "true"
+          effect = "NO_SCHEDULE"
+        }
+      }
+    }
+
+    # This node group is for core addons such as CoreDNS
+    default = {
+      instance_types = ["m5.large"]
+
+      min_size     = 1
+      max_size     = 2
+      desired_size = 2
+    }
+  }
+
+  tags = local.tags
+}

patterns/aws-neuron-efa/helm.tf

@@ -0,0 +1,50 @@
+data "aws_ecrpublic_authorization_token" "token" {
+  provider = aws.ecr
+}
+
+################################################################################
+# Helm charts
+################################################################################
+
+resource "helm_release" "neuron" {
+  name             = "neuron"
+  repository       = "oci://public.ecr.aws/neuron"
+  chart            = "neuron-helm-chart"
+  version          = "1.0.0"
+  namespace        = "neuron"
+  create_namespace = true
+  wait             = false
+
+  # Public ECR
+  repository_username = data.aws_ecrpublic_authorization_token.token.user_name
+  repository_password = data.aws_ecrpublic_authorization_token.token.password
+
+  values = [
+    <<-EOT
+      nodeSelector:
+        aws.amazon.com/neuron.present: 'true'
+      npd:
+        enabled: false
+    EOT
+  ]
+}
+
+resource "helm_release" "aws_efa_device_plugin" {
+  name       = "aws-efa-k8s-device-plugin"
+  repository = "https://aws.github.io/eks-charts"
+  chart      = "aws-efa-k8s-device-plugin"
+  version    = "v0.5.5"
+  namespace  = "kube-system"
+  wait       = false
+
+  values = [
+    <<-EOT
+      nodeSelector:
+        vpc.amazonaws.com/efa.present: 'true'
+      tolerations:
+        - key: aws.amazon.com/neuron
+          operator: Exists
+          effect: NoSchedule
+    EOT
+  ]
+}

patterns/aws-neuron-efa/main.tf

@@ -0,0 +1,109 @@
+terraform {
+  required_version = ">= 1.3"
+
+  required_providers {
+    aws = {
+      source  = "hashicorp/aws"
+      version = ">= 5.70"
+    }
+    helm = {
+      source  = "hashicorp/helm"
+      version = ">= 2.16"
+    }
+  }
+
+  # ##  Used for end-to-end testing on project; update to suit your needs
+  # backend "s3" {
+  #   bucket = "terraform-ssp-github-actions-state"
+  #   region = "us-west-2"
+  #   key    = "e2e/aws-neuron-efa/terraform.tfstate"
+  # }
+}
+
+provider "aws" {
+  region = local.region
+}
+
+# This provider is required for Public ECR. Public ECR is only available in us-east-1
+# If your region is same as us-east-1 then you can just use one aws provider
+provider "aws" {
+  alias  = "ecr"
+  region = "us-east-1"
+}
+
+provider "helm" {
+  kubernetes {
+    host                   = module.eks.cluster_endpoint
+    cluster_ca_certificate = base64decode(module.eks.cluster_certificate_authority_data)
+
+    exec {
+      api_version = "client.authentication.k8s.io/v1beta1"
+      command     = "aws"
+      # This requires the awscli to be installed locally where Terraform is executed
+      args = ["eks", "get-token", "--cluster-name", module.eks.cluster_name]
+    }
+  }
+}
+
+################################################################################
+# Common data/locals
+################################################################################
+
+data "aws_availability_zones" "available" {
+  # Do not include local zones
+  filter {
+    name   = "opt-in-status"
+    values = ["opt-in-not-required"]
+  }
+}
+
+locals {
+  name   = basename(path.cwd)
+  region = "us-east-2"
+
+  vpc_cidr = "10.0.0.0/16"
+  azs      = slice(data.aws_availability_zones.available.names, 0, 3)
+
+  tags = {
+    Blueprint  = local.name
+    GithubRepo = "github.com/aws-ia/terraform-aws-eks-blueprints"
+  }
+}
+
+################################################################################
+# Output
+################################################################################
+
+output "configure_kubectl" {
+  description = "Configure kubectl: make sure you're logged in with the correct AWS profile and run the following command to update your kubeconfig"
+  value       = "aws eks --region ${local.region} update-kubeconfig --name ${module.eks.cluster_name}"
+}
+
+################################################################################
+# Supporting Resources
+################################################################################
+
+module "vpc" {
+  source  = "terraform-aws-modules/vpc/aws"
+  version = "~> 5.0"
+
+  name = local.name
+  cidr = local.vpc_cidr
+
+  azs             = local.azs
+  private_subnets = [for k, v in local.azs : cidrsubnet(local.vpc_cidr, 4, k)]
+  public_subnets  = [for k, v in local.azs : cidrsubnet(local.vpc_cidr, 8, k + 48)]
+
+  enable_nat_gateway = true
+  single_nat_gateway = true
+
+  public_subnet_tags = {
+    "kubernetes.io/role/elb" = 1
+  }
+
+  private_subnet_tags = {
+    "kubernetes.io/role/internal-elb" = 1
+  }
+
+  tags = local.tags
+}