Skip to content

Latest commit

 

History

History
856 lines (689 loc) · 21.2 KB

File metadata and controls

856 lines (689 loc) · 21.2 KB

Quick Start Guide: AWS Neuron on EKS

Complete deployment guide for Trainium (training), Inferentia2 (inference), and vLLM-Neuron on Amazon EKS.

Time to complete: 30-45 minutes per deployment path


Prerequisites

Required Tools

# Check versions
aws --version          # AWS CLI 2.13+
kubectl version        # kubectl 1.28+
eksctl version         # eksctl 0.165+
helm version           # helm 3.12+

EKS Cluster Requirements

  • EKS version: 1.28 or higher
  • Karpenter: 0.32+ installed
  • IRSA: Enabled for service accounts
  • VPC: Subnet tags for Karpenter discovery
  • IAM: Karpenter node role with Neuron permissions

AWS Permissions

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Effect": "Allow",
      "Action": [
        "ec2:RunInstances",
        "ec2:CreateFleet",
        "ec2:CreateTags",
        "ec2:DescribeInstances",
        "ec2:DescribeInstanceTypes",
        "ec2:DescribeImages",
        "pricing:GetProducts",
        "ssm:GetParameter"
      ],
      "Resource": "*"
    }
  ]
}

Deployment Path Selection

Choose based on your use case:

Path Use Case Time Complexity
Path 1 Production LLM API serving 20 min ⭐⭐☆☆☆
Path 2 LLM training/fine-tuning 30 min ⭐⭐⭐☆☆
Path 3 High-throughput inference 25 min ⭐⭐⭐☆☆
Path 4 Complete MLOps pipeline 45 min ⭐⭐⭐⭐☆

Path 1: Inferentia2 for LLM Inference

What you'll deploy: Llama 3 8B inference on Inf2 with Neuron SDK

Cost: ~$0.76/hour (Inf2.xlarge)

Step 1: Install Neuron Device Plugin (5 minutes)

The Neuron Device Plugin exposes Inferentia2/Trainium as Kubernetes resources.

# Deploy device plugin
kubectl apply -f https://raw.githubusercontent.com/aws-neuron/aws-neuron-sdk/main/src/k8s/k8s-neuron-device-plugin.yml

# Verify DaemonSet is running
kubectl get daemonset neuron-device-plugin-daemonset -n kube-system

# Expected output:
# NAME                              DESIRED   CURRENT   READY   UP-TO-DATE   AVAILABLE
# neuron-device-plugin-daemonset    3         3         3       3            3

What this does: Registers aws.amazon.com/neuroncore and aws.amazon.com/neurondevice as allocatable resources.

Step 2: Create Karpenter NodePool for Inf2 (2 minutes)

# Apply Inf2 NodePool
kubectl apply -f - <<EOF
apiVersion: karpenter.sh/v1beta1
kind: NodePool
metadata:
  name: inferentia2-pool
spec:
  template:
    spec:
      requirements:
        - key: karpenter.sh/capacity-type
          operator: In
          values: ["on-demand"]
        - key: node.kubernetes.io/instance-type
          operator: In
          values:
            - inf2.xlarge
            - inf2.8xlarge
            - inf2.24xlarge
            - inf2.48xlarge
        - key: kubernetes.io/arch
          operator: In
          values: ["amd64"]
      
      nodeClassRef:
        name: neuron-nodes
  
  limits:
    cpu: "256"
    memory: 1024Gi
  
  disruption:
    consolidationPolicy: WhenUnderutilized
    expireAfter: 720h
---
apiVersion: karpenter.k8s.aws/v1beta1
kind: EC2NodeClass
metadata:
  name: neuron-nodes
spec:
  amiFamily: AL2
  role: KarpenterNodeRole-${CLUSTER_NAME}
  
  subnetSelectorTerms:
    - tags:
        karpenter.sh/discovery: ${CLUSTER_NAME}
  
  securityGroupSelectorTerms:
    - tags:
        karpenter.sh/discovery: ${CLUSTER_NAME}
  
  userData: |
    #!/bin/bash
    # Install Neuron driver and runtime
    . /etc/os-release
    sudo tee /etc/yum.repos.d/neuron.repo > /dev/null <<NEURON_REPO
    [neuron]
    name=Neuron YUM Repository
    baseurl=https://yum.repos.neuron.amazonaws.com
    enabled=1
    metadata_expire=0
    NEURON_REPO
    
    sudo rpm --import https://yum.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB
    sudo yum install -y aws-neuronx-dkms aws-neuronx-tools-2.18.1.0
    
    # Enable Neuron runtime
    export PATH=/opt/aws/neuron/bin:$PATH
    
    # Configure containerd for Neuron
    sudo mkdir -p /etc/containerd/
    sudo tee /etc/containerd/config.toml > /dev/null <<CONTAINERD_CONFIG
    version = 2
    [plugins."io.containerd.grpc.v1.cri".containerd.runtimes.neuron]
      runtime_type = "io.containerd.runc.v2"
      [plugins."io.containerd.grpc.v1.cri".containerd.runtimes.neuron.options]
        BinaryName = "/usr/bin/runc"
    CONTAINERD_CONFIG
    
    sudo systemctl restart containerd
EOF

Important: Replace ${CLUSTER_NAME} with your actual cluster name.

Step 3: Deploy Llama 3 8B Inference (5 minutes)

# Create deployment
kubectl apply -f - <<EOF
apiVersion: apps/v1
kind: Deployment
metadata:
  name: llama3-inf2
  labels:
    app: llama3-inf2
spec:
  replicas: 1
  selector:
    matchLabels:
      app: llama3-inf2
  template:
    metadata:
      labels:
        app: llama3-inf2
    spec:
      containers:
        - name: neuron-inference
          image: public.ecr.aws/neuron/pytorch-inference-neuronx:2.1.2-neuronx-py310-sdk2.18.1-ubuntu20.04
          
          command:
            - python3
            - -c
            - |
              from transformers_neuronx import NeuronAutoModelForCausalLM
              from transformers import AutoTokenizer
              import torch
              
              model_id = "meta-llama/Meta-Llama-3-8B"
              
              # Load model for Neuron
              model = NeuronAutoModelForCausalLM.from_pretrained(
                  model_id,
                  torch_dtype=torch.bfloat16,
                  low_cpu_mem_usage=True
              )
              model.to_neuron()
              
              tokenizer = AutoTokenizer.from_pretrained(model_id)
              
              # Simple HTTP server
              from flask import Flask, request, jsonify
              app = Flask(__name__)
              
              @app.route('/v1/completions', methods=['POST'])
              def complete():
                  data = request.json
                  prompt = data.get('prompt', '')
                  max_tokens = data.get('max_tokens', 100)
                  
                  inputs = tokenizer(prompt, return_tensors='pt')
                  outputs = model.generate(
                      **inputs,
                      max_new_tokens=max_tokens,
                      do_sample=True,
                      temperature=0.7
                  )
                  
                  text = tokenizer.decode(outputs[0], skip_special_tokens=True)
                  
                  return jsonify({
                      'id': 'neuron-' + str(hash(prompt)),
                      'object': 'text_completion',
                      'model': model_id,
                      'choices': [{
                          'text': text,
                          'index': 0,
                          'finish_reason': 'stop'
                      }]
                  })
              
              app.run(host='0.0.0.0', port=8000)
          
          ports:
            - containerPort: 8000
              name: http
          
          resources:
            requests:
              aws.amazon.com/neuroncore: "2"
              memory: 16Gi
              cpu: "4"
            limits:
              aws.amazon.com/neuroncore: "2"
              memory: 24Gi
              cpu: "8"
          
          env:
            - name: NEURON_RT_NUM_CORES
              value: "2"
            - name: NEURON_RT_VISIBLE_CORES
              value: "0-1"
          
          livenessProbe:
            httpGet:
              path: /health
              port: 8000
            initialDelaySeconds: 300
            periodSeconds: 30
          
          readinessProbe:
            httpGet:
              path: /health
              port: 8000
            initialDelaySeconds: 240
            periodSeconds: 10
---
apiVersion: v1
kind: Service
metadata:
  name: llama3-inf2
spec:
  type: LoadBalancer
  ports:
    - port: 80
      targetPort: 8000
      protocol: TCP
      name: http
  selector:
    app: llama3-inf2
EOF

Step 4: Wait for Deployment (8 minutes)

# Watch pod status
kubectl get pods -l app=llama3-inf2 -w

# Wait for READY 1/1
# This takes 5-8 minutes (node provision + model load)

What's happening:

  1. Karpenter provisions Inf2.xlarge node (2-3 min)
  2. Neuron driver initializes (1 min)
  3. Model downloads from Hugging Face (2-3 min)
  4. Model compiles for Neuron (2-3 min)

Step 5: Test Inference (2 minutes)

# Get service endpoint
export ENDPOINT=$(kubectl get svc llama3-inf2 -o jsonpath='{.status.loadBalancer.ingress[0].hostname}')

# Test completion
curl http://${ENDPOINT}/v1/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "meta-llama/Meta-Llama-3-8B",
    "prompt": "AWS Neuron is",
    "max_tokens": 50
  }'

Expected output:

{
  "id": "neuron-123456789",
  "object": "text_completion",
  "model": "meta-llama/Meta-Llama-3-8B",
  "choices": [{
    "text": "AWS Neuron is a software development kit (SDK) that enables developers to deploy machine learning models on AWS Trainium and Inferentia accelerators...",
    "index": 0,
    "finish_reason": "stop"
  }]
}

Step 6: Validate Performance (3 minutes)

# Run benchmark
./scripts/benchmark-inference.sh \
  --endpoint http://${ENDPOINT} \
  --model llama3-8b \
  --requests 100 \
  --concurrency 10

# Expected results:
# Throughput: 1,100 tokens/s
# P50 latency: 65ms
# P99 latency: 120ms

Cleanup

kubectl delete deployment llama3-inf2
kubectl delete service llama3-inf2
kubectl delete nodepool inferentia2-pool

Path 2: Trainium for LLM Training

What you'll deploy: Llama 2 7B fine-tuning on Trn1

Cost: ~$24.78/hour (Trn1.32xlarge)

Step 1: Install Neuron Device Plugin (if not done)

kubectl apply -f https://raw.githubusercontent.com/aws-neuron/aws-neuron-sdk/main/src/k8s/k8s-neuron-device-plugin.yml

Step 2: Create Karpenter NodePool for Trn1 (2 minutes)

kubectl apply -f - <<EOF
apiVersion: karpenter.sh/v1beta1
kind: NodePool
metadata:
  name: trainium-pool
spec:
  template:
    spec:
      requirements:
        - key: karpenter.sh/capacity-type
          operator: In
          values: ["on-demand"]
        - key: node.kubernetes.io/instance-type
          operator: In
          values:
            - trn1.2xlarge
            - trn1.32xlarge
            - trn1n.32xlarge
        - key: kubernetes.io/arch
          operator: In
          values: ["amd64"]
      
      nodeClassRef:
        name: neuron-nodes
      
      taints:
        - key: aws.amazon.com/neuron
          value: "true"
          effect: NoSchedule
  
  limits:
    cpu: "512"
    memory: 2048Gi
EOF

Step 3: Create Training Job (5 minutes)

kubectl apply -f - <<EOF
apiVersion: batch/v1
kind: Job
metadata:
  name: llama2-7b-training
spec:
  template:
    spec:
      restartPolicy: OnFailure
      
      tolerations:
        - key: aws.amazon.com/neuron
          operator: Equal
          value: "true"
          effect: NoSchedule
      
      containers:
        - name: training
          image: public.ecr.aws/neuron/pytorch-training-neuronx:2.1.2-neuronx-py310-sdk2.18.1-ubuntu20.04
          
          command:
            - torchrun
            - --nproc_per_node=32
            - train.py
            - --model_id=meta-llama/Llama-2-7b-hf
            - --dataset=wikitext
            - --batch_size=8
            - --epochs=3
            - --learning_rate=5e-5
          
          volumeMounts:
            - name: training-script
              mountPath: /workspace/train.py
              subPath: train.py
            - name: model-output
              mountPath: /output
          
          resources:
            requests:
              aws.amazon.com/neuroncore: "32"
              memory: 256Gi
              cpu: "64"
            limits:
              aws.amazon.com/neuroncore: "32"
              memory: 384Gi
              cpu: "96"
          
          env:
            - name: NEURON_RT_NUM_CORES
              value: "32"
            - name: NEURON_RT_VISIBLE_CORES
              value: "0-31"
            - name: NEURON_CC_FLAGS
              value: "--model-type=transformer --auto-cast=bf16"
      
      volumes:
        - name: training-script
          configMap:
            name: training-script
        - name: model-output
          persistentVolumeClaim:
            claimName: model-output-pvc
EOF

Step 4: Monitor Training (15 minutes)

# Watch job status
kubectl get job llama2-7b-training -w

# Stream logs
kubectl logs -f job/llama2-7b-training

# Monitor Neuron metrics
kubectl exec -it $(kubectl get pod -l job-name=llama2-7b-training -o name) -- neuron-monitor

Expected output:

Epoch 1/3: loss=2.134, throughput=120K tokens/s
Epoch 2/3: loss=1.892, throughput=118K tokens/s
Epoch 3/3: loss=1.756, throughput=121K tokens/s
Training complete. Model saved to /output/llama2-7b-finetuned/

Step 5: Export Model (2 minutes)

# Copy model to S3
kubectl exec -it $(kubectl get pod -l job-name=llama2-7b-training -o name) -- \
  aws s3 sync /output/llama2-7b-finetuned/ s3://your-bucket/models/llama2-7b-finetuned/

Path 3: vLLM-Neuron (High-Performance Inference)

What you'll deploy: vLLM with PagedAttention on Inferentia2

Cost: ~$0.76/hour (Inf2.xlarge)

Step 1: Install Prerequisites

# Neuron device plugin (if not done)
kubectl apply -f https://raw.githubusercontent.com/aws-neuron/aws-neuron-sdk/main/src/k8s/k8s-neuron-device-plugin.yml

# Apply Inf2 NodePool from Path 1

Step 2: Deploy vLLM Server (5 minutes)

kubectl apply -f deployments/vllm/vllm-llama3-inf2.yaml

# File contents:
# - vLLM container with Neuron backend
# - PagedAttention enabled
# - Multi-model serving support
# - OpenAI-compatible API

Step 3: Test vLLM Endpoints (2 minutes)

# Get endpoint
export VLLM_ENDPOINT=$(kubectl get svc vllm-llama3 -o jsonpath='{.status.loadBalancer.ingress[0].hostname}')

# Test completion
curl http://${VLLM_ENDPOINT}/v1/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "meta-llama/Meta-Llama-3-8B",
    "prompt": "Explain vLLM in one sentence:",
    "max_tokens": 50,
    "temperature": 0.7
  }'

# Test chat completion
curl http://${VLLM_ENDPOINT}/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "meta-llama/Meta-Llama-3-8B",
    "messages": [
      {"role": "user", "content": "What is AWS Neuron?"}
    ],
    "max_tokens": 100
  }'

Step 4: Enable Token Streaming (1 minute)

# Streaming completion
curl http://${VLLM_ENDPOINT}/v1/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "meta-llama/Meta-Llama-3-8B",
    "prompt": "Write a story about AI:",
    "max_tokens": 200,
    "stream": true
  }'

What you'll see: Tokens streamed in real-time (SSE format)

Step 5: Benchmark vLLM Performance (5 minutes)

./scripts/benchmark-vllm.sh \
  --endpoint http://${VLLM_ENDPOINT} \
  --model llama3-8b \
  --requests 1000 \
  --concurrency 50

# Expected results with vLLM-Neuron:
# Throughput: 3,200 tokens/s (vs 1,100 without vLLM)
# P50 latency: 45ms
# P99 latency: 95ms
# GPU memory efficiency: 85%+

Path 4: Hybrid (Training + Inference Pipeline)

What you'll deploy: Complete MLOps pipeline with Trn1 training and Inf2 inference

Time: 45 minutes
Cost: $24.78/hour (training) + $0.76/hour (inference)

Architecture

1. Fine-tune on Trn1.32xlarge
2. Save model to EFS
3. Deploy inference on Inf2.xlarge
4. Monitor end-to-end pipeline

Step 1: Create Shared EFS Storage (5 minutes)

# Create EFS filesystem
aws efs create-file-system \
  --performance-mode generalPurpose \
  --throughput-mode bursting \
  --encrypted \
  --tags Key=Name,Value=neuron-models

export EFS_ID=$(aws efs describe-file-systems --query 'FileSystems[?Tags[?Key==`Name` && Value==`neuron-models`]].FileSystemId' --output text)

# Create mount targets
for SUBNET in $(aws ec2 describe-subnets --filters "Name=tag:karpenter.sh/discovery,Values=${CLUSTER_NAME}" --query 'Subnets[].SubnetId' --output text); do
  aws efs create-mount-target \
    --file-system-id $EFS_ID \
    --subnet-id $SUBNET \
    --security-groups $(aws eks describe-cluster --name ${CLUSTER_NAME} --query 'cluster.resourcesVpcConfig.clusterSecurityGroupId' --output text)
done

# Install EFS CSI driver
helm repo add aws-efs-csi-driver https://kubernetes-sigs.github.io/aws-efs-csi-driver/
helm install aws-efs-csi-driver aws-efs-csi-driver/aws-efs-csi-driver \
  --namespace kube-system

# Create PVC
kubectl apply -f - <<EOF
apiVersion: v1
kind: PersistentVolume
metadata:
  name: efs-models-pv
spec:
  capacity:
    storage: 100Gi
  accessModes:
    - ReadWriteMany
  csi:
    driver: efs.csi.aws.com
    volumeHandle: ${EFS_ID}
---
apiVersion: v1
kind: PersistentVolumeClaim
metadata:
  name: efs-models-pvc
spec:
  accessModes:
    - ReadWriteMany
  resources:
    requests:
      storage: 100Gi
  volumeName: efs-models-pv
EOF

Step 2: Run Training Job with EFS (15 minutes)

# Deploy training job (from Path 2) with EFS mount
kubectl apply -f deployments/trainium/llama-training-efs.yaml

# Monitor
kubectl logs -f job/llama2-training

Step 3: Deploy Inference from EFS (10 minutes)

# Deploy inference pointing to EFS model
kubectl apply -f deployments/inferentia2/llama-inference-efs.yaml

# Wait for ready
kubectl wait --for=condition=ready pod -l app=llama-inference --timeout=600s

Step 4: Test End-to-End Pipeline (5 minutes)

# Get inference endpoint
export INFERENCE_ENDPOINT=$(kubectl get svc llama-inference -o jsonpath='{.status.loadBalancer.ingress[0].hostname}')

# Test fine-tuned model
curl http://${INFERENCE_ENDPOINT}/v1/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "llama2-7b-finetuned",
    "prompt": "Test prompt for fine-tuned model:",
    "max_tokens": 100
  }'

Step 5: Set Up Monitoring (10 minutes)

# Deploy Neuron Monitor DaemonSet
kubectl apply -f monitoring/neuron-monitor-daemonset.yaml

# Deploy Prometheus rules
kubectl apply -f monitoring/prometheus-rules.yaml

# Import Grafana dashboard
kubectl apply -f monitoring/grafana-dashboard.json

Grafana metrics available:

  • NeuronCore utilization (training + inference)
  • Model throughput (tokens/s)
  • Training loss curves
  • Inference latency (P50, P95, P99)
  • Cost per token

Validation Checklist

After deployment, verify:

Neuron Device Plugin

kubectl get daemonset -n kube-system neuron-device-plugin-daemonset
# STATUS: Running on all Inf2/Trn1 nodes

Neuron Resources Advertised

kubectl describe node <inf2-or-trn1-node> | grep neuron
# Should show:
#   aws.amazon.com/neuroncore: 2 (or 32 for Trn1)
#   aws.amazon.com/neurondevice: 1

Model Loading

kubectl logs <pod-name> | grep "Model loaded"
# Should show successful model load

Inference Working

curl http://<endpoint>/v1/completions -X POST -d '{"prompt": "test"}'
# Should return valid completion

Neuron Monitor

kubectl exec -it <pod-name> -- neuron-top
# Should show NeuronCore utilization

Troubleshooting

Issue 1: Neuron Device Plugin Not Running

Symptom:

kubectl get daemonset -n kube-system neuron-device-plugin-daemonset
# DESIRED: 3, READY: 0

Solution:

# Check logs
kubectl logs -n kube-system -l app.kubernetes.io/name=neuron-device-plugin

# Common fix: Update device plugin version
kubectl delete daemonset -n kube-system neuron-device-plugin-daemonset
kubectl apply -f https://raw.githubusercontent.com/aws-neuron/aws-neuron-sdk/main/src/k8s/k8s-neuron-device-plugin.yml

Issue 2: Model Loading Fails

Symptom:

Error: Failed to load model: No Neuron devices found

Solution:

# Verify Neuron driver installed
kubectl exec -it <pod-name> -- ls /opt/aws/neuron/

# Reinstall if missing (via userData in EC2NodeClass)

Issue 3: Low Throughput

Symptom: Getting <500 tokens/s on Inf2.xlarge (expected: 1,100+)

Solution:

# Check NeuronCore allocation
kubectl describe pod <pod-name> | grep neuroncore
# Should request 2+ cores

# Enable batch inference
# Update deployment with:
#   --max-batch-size=16
#   --max-batch-wait-time=50ms

See docs/troubleshooting.md for complete guide.


Next Steps

Optimize Performance

Scale Production

Monitor & Debug


Questions? See docs/ or open an issue.