Skip to content

Commit 04a1b4b

Browse files
bryantbiggsbryantbiggs
authored
feat: Add pattern for Neuron using Trainium and new AL2023 Neuron AMI (#2034)
1 parent 9ec1d47 commit 04a1b4b

File tree

5 files changed

+313
-0
lines changed

5 files changed

+313
-0
lines changed

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

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
---
2+
title: AWS Neuron with EFA
3+
---
4+
5+
{%
6+
include-markdown "../../../patterns/aws-neuron-efa/README.md"
7+
%}

‎patterns/aws-neuron-efa/README.md

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
# EKS Cluster w/ AWS Neuron Devices and EFA for Machine Learning
2+
3+
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.
4+
5+
The following components are demonstrated in this pattern:
6+
7+
- 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.
8+
- A node group of `trn1.32xlarge` instances with:
9+
- all x8 [EFA network interfaces](https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/efa.html) enabled
10+
- 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
11+
- a common taint of `"aws.amazon.com/neuron:NoSchedule"` to ensure only the intended applications are permitted to run on the nodes created
12+
- two labels identifying that this nodegroup supports AWS Neuron and EFA devices; allowing pods to use node selectors with these labels
13+
- 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
14+
- 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
15+
- 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
16+
- 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.
17+
18+
## Code
19+
20+
```terraform hl_lines="26-28 34-80"
21+
{% include "../../patterns/aws-neuron-efa/eks.tf" %}
22+
```
23+
24+
```terraform hl_lines="9-50"
25+
{% include "../../patterns/aws-neuron-efa/helm.tf" %}
26+
```
27+
28+
## Deploy
29+
30+
See [here](https://aws-ia.github.io/terraform-aws-eks-blueprints/getting-started/#prerequisites) for the prerequisites and steps to deploy this pattern.
31+
32+
## Validate
33+
34+
1. List the nodes and their instance type:
35+
36+
```sh
37+
kubectl get nodes -L node.kubernetes.io/instance-type
38+
```
39+
40+
```text
41+
NAME STATUS ROLES AGE VERSION INSTANCE-TYPE
42+
ip-10-0-12-200.us-east-2.compute.internal Ready <none> 82m v1.31.0-eks-a737599 m5.large
43+
ip-10-0-24-248.us-east-2.compute.internal Ready <none> 82m v1.31.0-eks-a737599 m5.large
44+
ip-10-0-39-213.us-east-2.compute.internal Ready <none> 75m v1.31.0-eks-a737599 trn1.32xlarge
45+
ip-10-0-43-172.us-east-2.compute.internal Ready <none> 75m v1.31.0-eks-a737599 trn1.32xlarge
46+
```
47+
48+
You should see two EFA-enabled (in this example `trn1.32xlarge`) nodes in the list.
49+
50+
## Destroy
51+
52+
{%
53+
include-markdown "../../docs/_partials/destroy.md"
54+
%}

‎patterns/aws-neuron-efa/eks.tf

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
################################################################################
2+
# Cluster
3+
################################################################################
4+
5+
module "eks" {
6+
source = "terraform-aws-modules/eks/aws"
7+
version = "~> 20.26"
8+
9+
cluster_name = local.name
10+
cluster_version = "1.31"
11+
12+
# Give the Terraform identity admin access to the cluster
13+
# which will allow it to deploy resources into the cluster
14+
enable_cluster_creator_admin_permissions = true
15+
cluster_endpoint_public_access = true
16+
17+
cluster_addons = {
18+
coredns = {}
19+
eks-pod-identity-agent = {}
20+
kube-proxy = {}
21+
vpc-cni = {
22+
most_recent = true
23+
}
24+
}
25+
26+
# Add security group rules on the node group security group to
27+
# allow EFA traffic
28+
enable_efa_support = true
29+
30+
vpc_id = module.vpc.vpc_id
31+
subnet_ids = module.vpc.private_subnets
32+
33+
eks_managed_node_groups = {
34+
neuron-efa = {
35+
# The EKS AL2023 Neuron AMI provides all of the necessary components
36+
# for accelerated workloads w/ EFA
37+
ami_type = "AL2023_x86_64_NEURON"
38+
instance_types = ["trn1.32xlarge"]
39+
40+
# Mount instance store volumes in RAID-0 for kubelet and containerd
41+
# https://github.com/awslabs/amazon-eks-ami/blob/master/doc/USER_GUIDE.md#raid-0-for-kubelet-and-containerd-raid0
42+
cloudinit_pre_nodeadm = [
43+
{
44+
content_type = "application/node.eks.aws"
45+
content = <<-EOT
46+
---
47+
apiVersion: node.eks.aws/v1alpha1
48+
kind: NodeConfig
49+
spec:
50+
instance:
51+
localStorage:
52+
strategy: RAID0
53+
EOT
54+
}
55+
]
56+
57+
min_size = 2
58+
max_size = 2
59+
desired_size = 2
60+
61+
# This will:
62+
# 1. Create a placement group to place the instances close to one another
63+
# 2. Ignore subnets that reside in AZs that do not support the instance type
64+
# 3. Expose all of the available EFA interfaces on the launch template
65+
enable_efa_support = true
66+
67+
labels = {
68+
"vpc.amazonaws.com/efa.present" = "true"
69+
"aws.amazon.com/neuron.present" = "true"
70+
}
71+
72+
taints = {
73+
# Ensure only Neuron workloads are scheduled on this node group
74+
gpu = {
75+
key = "aws.amazon.com/neuron"
76+
value = "true"
77+
effect = "NO_SCHEDULE"
78+
}
79+
}
80+
}
81+
82+
# This node group is for core addons such as CoreDNS
83+
default = {
84+
instance_types = ["m5.large"]
85+
86+
min_size = 1
87+
max_size = 2
88+
desired_size = 2
89+
}
90+
}
91+
92+
tags = local.tags
93+
}

‎patterns/aws-neuron-efa/helm.tf

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
data "aws_ecrpublic_authorization_token" "token" {
2+
provider = aws.ecr
3+
}
4+
5+
################################################################################
6+
# Helm charts
7+
################################################################################
8+
9+
resource "helm_release" "neuron" {
10+
name = "neuron"
11+
repository = "oci://public.ecr.aws/neuron"
12+
chart = "neuron-helm-chart"
13+
version = "1.0.0"
14+
namespace = "neuron"
15+
create_namespace = true
16+
wait = false
17+
18+
# Public ECR
19+
repository_username = data.aws_ecrpublic_authorization_token.token.user_name
20+
repository_password = data.aws_ecrpublic_authorization_token.token.password
21+
22+
values = [
23+
<<-EOT
24+
nodeSelector:
25+
aws.amazon.com/neuron.present: 'true'
26+
npd:
27+
enabled: false
28+
EOT
29+
]
30+
}
31+
32+
resource "helm_release" "aws_efa_device_plugin" {
33+
name = "aws-efa-k8s-device-plugin"
34+
repository = "https://aws.github.io/eks-charts"
35+
chart = "aws-efa-k8s-device-plugin"
36+
version = "v0.5.5"
37+
namespace = "kube-system"
38+
wait = false
39+
40+
values = [
41+
<<-EOT
42+
nodeSelector:
43+
vpc.amazonaws.com/efa.present: 'true'
44+
tolerations:
45+
- key: aws.amazon.com/neuron
46+
operator: Exists
47+
effect: NoSchedule
48+
EOT
49+
]
50+
}

‎patterns/aws-neuron-efa/main.tf

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
terraform {
2+
required_version = ">= 1.3"
3+
4+
required_providers {
5+
aws = {
6+
source = "hashicorp/aws"
7+
version = ">= 5.70"
8+
}
9+
helm = {
10+
source = "hashicorp/helm"
11+
version = ">= 2.16"
12+
}
13+
}
14+
15+
# ## Used for end-to-end testing on project; update to suit your needs
16+
# backend "s3" {
17+
# bucket = "terraform-ssp-github-actions-state"
18+
# region = "us-west-2"
19+
# key = "e2e/aws-neuron-efa/terraform.tfstate"
20+
# }
21+
}
22+
23+
provider "aws" {
24+
region = local.region
25+
}
26+
27+
# This provider is required for Public ECR. Public ECR is only available in us-east-1
28+
# If your region is same as us-east-1 then you can just use one aws provider
29+
provider "aws" {
30+
alias = "ecr"
31+
region = "us-east-1"
32+
}
33+
34+
provider "helm" {
35+
kubernetes {
36+
host = module.eks.cluster_endpoint
37+
cluster_ca_certificate = base64decode(module.eks.cluster_certificate_authority_data)
38+
39+
exec {
40+
api_version = "client.authentication.k8s.io/v1beta1"
41+
command = "aws"
42+
# This requires the awscli to be installed locally where Terraform is executed
43+
args = ["eks", "get-token", "--cluster-name", module.eks.cluster_name]
44+
}
45+
}
46+
}
47+
48+
################################################################################
49+
# Common data/locals
50+
################################################################################
51+
52+
data "aws_availability_zones" "available" {
53+
# Do not include local zones
54+
filter {
55+
name = "opt-in-status"
56+
values = ["opt-in-not-required"]
57+
}
58+
}
59+
60+
locals {
61+
name = basename(path.cwd)
62+
region = "us-east-2"
63+
64+
vpc_cidr = "10.0.0.0/16"
65+
azs = slice(data.aws_availability_zones.available.names, 0, 3)
66+
67+
tags = {
68+
Blueprint = local.name
69+
GithubRepo = "github.com/aws-ia/terraform-aws-eks-blueprints"
70+
}
71+
}
72+
73+
################################################################################
74+
# Output
75+
################################################################################
76+
77+
output "configure_kubectl" {
78+
description = "Configure kubectl: make sure you're logged in with the correct AWS profile and run the following command to update your kubeconfig"
79+
value = "aws eks --region ${local.region} update-kubeconfig --name ${module.eks.cluster_name}"
80+
}
81+
82+
################################################################################
83+
# Supporting Resources
84+
################################################################################
85+
86+
module "vpc" {
87+
source = "terraform-aws-modules/vpc/aws"
88+
version = "~> 5.0"
89+
90+
name = local.name
91+
cidr = local.vpc_cidr
92+
93+
azs = local.azs
94+
private_subnets = [for k, v in local.azs : cidrsubnet(local.vpc_cidr, 4, k)]
95+
public_subnets = [for k, v in local.azs : cidrsubnet(local.vpc_cidr, 8, k + 48)]
96+
97+
enable_nat_gateway = true
98+
single_nat_gateway = true
99+
100+
public_subnet_tags = {
101+
"kubernetes.io/role/elb" = 1
102+
}
103+
104+
private_subnet_tags = {
105+
"kubernetes.io/role/internal-elb" = 1
106+
}
107+
108+
tags = local.tags
109+
}

0 commit comments

Comments
 (0)