Kubernetes environment setup for Neuron
Contents
This document is relevant for: Inf1
, Inf2
, Trn1
, Trn1n
Kubernetes environment setup for Neuron#
Introduction#
Customers that use Kubernetes can conveniently integrate Inf1/Trn1 instances into their workflows. This tutorial will go through deploying the neuron device plugin daemonset and also how to allocate neuron cores or devices to application pods.
Prerequisite
Working kubernetes cluster
Inf1/Trn1 instances as worker nodes with attached roles allowing: * ECR read access policy to retrieve container images from ECR: arn:aws:iam::aws:policy/AmazonEC2ContainerRegistryReadOnly
Tutorial Docker environment setup: to install required packages in the worker nodes. With EKS, the EKS optimized accelarated AMI has the necessary neuron components installed
Kubernetes node object has instance-type set to inf1/trn1 types. For ex,
"node.kubernetes.io/instance-type": "inf1.2xlarge"
or"node.kubernetes.io/instance-type": "trn1.2xlarge"
Deploy Neuron Device Plugin
Neuron device plugin exposes Neuron cores & devices to kubernetes as a resource. aws.amazon.com/neuroncore, aws.amazon.com/neurondevice, aws.amazon.com/neuron are the resources that the neuron device plugin registers with the kubernetes. aws.amazon.com/neuroncore is used for allocating neuron cores to the container. aws.amazon.com/neurondevice is used for allocating neuron devices to the container. When neurondevice is used all the cores belonging to the device will be allocated to container. aws.amazon.com/neuron also allocates neurondevices and this exists just to be backward compatible with already existing installations. aws.amazon.com/neurondevice is the recommended resource for allocating devices to the container.
Make sure prequisite are satisified
Download the neuron device plugin yaml file.
k8s-neuron-device-plugin.yml
Download the neuron device plugin rbac yaml file. This enables permissions for device plugin to update the node and Pod annotations.
k8s-neuron-device-plugin-rbac.yml
Apply the Neuron device plugin as a daemonset on the cluster with the following command
kubectl apply -f k8s-neuron-device-plugin-rbac.yml kubectl apply -f k8s-neuron-device-plugin.yml
Verify that neuron device plugin is running
kubectl get ds neuron-device-plugin-daemonset --namespace kube-system
Expected result (with 2 nodes in cluster):
NAME DESIRED CURRENT READY UP-TO-DATE AVAILABLE NODE SELECTOR AGE neuron-device-plugin-daemonset 2 2 2 2 2 <none> 27h
Verify that the node has allocatable neuron cores and devices with the following command
kubectl get nodes "-o=custom-columns=NAME:.metadata.name,NeuronCore:.status.allocatable.aws\.amazon\.com/neuroncore"
Expected result:
NAME NeuronCore ip-192-168-65-41.us-west-2.compute.internal 32 ip-192-168-87-81.us-west-2.compute.internal 32
kubectl get nodes "-o=custom-columns=NAME:.metadata.name,NeuronDevice:.status.allocatable.aws\.amazon\.com/neurondevice"
Expected result:
NAME NeuronDevice ip-192-168-65-41.us-west-2.compute.internal 16 ip-192-168-87-81.us-west-2.compute.internal 16
Deploy Neuron Scheduler Extension
Neuron scheduler extension is required for scheduling pods that require more than one Neuron core or device resource. Refer Neuron Scheduler Extension Flow Diagram for details on how the neuron scheduler extension works. Neuron scheduler extension filter out nodes with non-contiguous core/device ids and enforces allocation of contiguous core/device ids for the PODs requiring it.
In cluster environments where there is no access to default scheduler, the neuron scheduler extension can be used with another scheduler. A new scheduler is added (along with the default scheduler) and then the pod’s that needs to run the neuron workload use this new scheduler. Neuron scheduler extension is added to this new scheduler. EKS natively does not yet support the neuron scheduler extension and so in the EKS environment this is the only way to add the neuron scheduler extension.
Make sure Neuron device plugin is running
Download the my scheduler
my-scheduler.yml
Download the scheduler extension
k8s-neuron-scheduler-eks.yml
Apply the neuron-scheduler-extension
kubectl apply -f k8s-neuron-scheduler-eks.yml
Apply the my-scheduler.yml
kubectl apply -f my-scheduler.yml
Check there are no errors in the my-scheduler pod logs and the k8s-neuron-scheduler pod is bound to a node
kubectl logs -n kube-system my-scheduler-79bd4cb788-hq2sq
I1012 15:30:21.629611 1 scheduler.go:604] "Successfully bound pod to node" pod="kube-system/k8s-neuron-scheduler-5d9d9d7988-xcpqm" node="ip-192-168-2-25.ec2.internal" evaluatedNodes=1 feasibleNodes=1
When running new pod’s that need to use the neuron scheduler extension, make sure it uses the my-scheduler as the scheduler. Sample pod spec is below
apiVersion: v1 kind: Pod metadata: name: <POD_NAME> spec: restartPolicy: Never schedulerName: my-scheduler containers: - name: <POD_NAME> command: ["<COMMAND>"] image: <IMAGE_NAME> resources: limits: cpu: "4" memory: 4Gi aws.amazon.com/neuroncore: 9 requests: cpu: "1" memory: 1Gi
Once the neuron workload pod is run, make sure logs in the k8s neuron scheduler has successfull filter/bind request
kubectl logs -n kube-system k8s-neuron-scheduler-5d9d9d7988-xcpqm
2022/10/12 15:41:16 POD nrt-test-5038 fits in Node:ip-192-168-2-25.ec2.internal 2022/10/12 15:41:16 Filtered nodes: [ip-192-168-2-25.ec2.internal] 2022/10/12 15:41:16 Failed nodes: map[] 2022/10/12 15:41:16 Finished Processing Filter Request...
2022/10/12 15:41:16 Executing Bind Request! 2022/10/12 15:41:16 Determine if the pod %v is NeuronDevice podnrt-test-5038 2022/10/12 15:41:16 Updating POD Annotation with alloc devices! 2022/10/12 15:41:16 Return aws.amazon.com/neuroncore 2022/10/12 15:41:16 neuronDevUsageMap for resource:aws.amazon.com/neuroncore in node: ip-192-168-2-25.ec2.internal is [false false false false false false false false false false false false false false false false] 2022/10/12 15:41:16 Allocated ids for POD nrt-test-5038 are: 0,1,2,3,4,5,6,7,8 2022/10/12 15:41:16 Try to bind pod nrt-test-5038 in default namespace to node ip-192-168-2-25.ec2.internal with &Binding{ObjectMeta:{nrt-test-5038 8da590b1-30bc-4335-b7e7-fe574f4f5538 0 0001-01-01 00:00:00 +0000 UTC <nil> <nil> map[] map[] [] [] []},Target:ObjectReference{Kind:Node,Namespace:,Name:ip-192-168-2-25.ec2.internal,UID:,APIVersion:,ResourceVersion:,FieldPath:,},} 2022/10/12 15:41:16 Updating the DevUsageMap since the bind is successful! 2022/10/12 15:41:16 Return aws.amazon.com/neuroncore 2022/10/12 15:41:16 neuronDevUsageMap for resource:aws.amazon.com/neuroncore in node: ip-192-168-2-25.ec2.internal is [false false false false false false false false false false false false false false false false] 2022/10/12 15:41:16 neuronDevUsageMap for resource:aws.amazon.com/neurondevice in node: ip-192-168-2-25.ec2.internal is [false false false false] 2022/10/12 15:41:16 Allocated devices list 0,1,2,3,4,5,6,7,8 for resource aws.amazon.com/neuroncore 2022/10/12 15:41:16 Allocated devices list [0] for other resource aws.amazon.com/neurondevice 2022/10/12 15:41:16 Allocated devices list [0] for other resource aws.amazon.com/neurondevice 2022/10/12 15:41:16 Allocated devices list [0] for other resource aws.amazon.com/neurondevice 2022/10/12 15:41:16 Allocated devices list [0] for other resource aws.amazon.com/neurondevice 2022/10/12 15:41:16 Allocated devices list [1] for other resource aws.amazon.com/neurondevice 2022/10/12 15:41:16 Allocated devices list [1] for other resource aws.amazon.com/neurondevice 2022/10/12 15:41:16 Allocated devices list [1] for other resource aws.amazon.com/neurondevice 2022/10/12 15:41:16 Allocated devices list [1] for other resource aws.amazon.com/neurondevice 2022/10/12 15:41:16 Allocated devices list [2] for other resource aws.amazon.com/neurondevice 2022/10/12 15:41:16 Return aws.amazon.com/neuroncore 2022/10/12 15:41:16 Succesfully updated the DevUsageMap [true true true true true true true true true false false false false false false false] and otherDevUsageMap [true true true false] after alloc for node ip-192-168-2-25.ec2.internal 2022/10/12 15:41:16 Finished executing Bind Request...
Make sure Neuron device plugin is running
Download the scheduler config map
k8s-neuron-scheduler-configmap.yml
Download the scheduler extension
k8s-neuron-scheduler.yml
Enable the kube-scheduler with option to use configMap for scheduler policy. In your cluster.yml Please update the spec section with the following
spec: kubeScheduler: usePolicyConfigMap: true
Launch the cluster
kops create -f cluster.yml kops create secret --name neuron-test-1.k8s.local sshpublickey admin -i ~/.ssh/id_rsa.pub kops update cluster --name neuron-test-1.k8s.local --yes
Apply the k8s-neuron-scheduler-configmap.yml [Registers neuron-scheduler-extension with kube-scheduler]
kubectl apply -f k8s-neuron-scheduler-configmap.yml
Launch the neuron-scheduler-extension
kubectl apply -f k8s-neuron-scheduler.yml
This document is relevant for: Inf1
, Inf2
, Trn1
, Trn1n