Step-by-Step: Deploy a Kubeflow 1.8 Pipeline for AI Training on Kubernetes 1.32 with AWS Graviton4

# Provider configuration for AWS and Kubernetes
# Pinned to AWS provider 5.83.0 to ensure EKS 1.32 support
terraform {
  required_providers {
    aws = {
      source  = "hashicorp/aws"
      version = "~> 5.83.0"
    }
    kubernetes = {
      source  = "hashicorp/kubernetes"
      version = "~> 2.23.0"
    }
  }
  required_version = ">= 1.7.0"
}

# Validate AWS region to prevent unsupported region errors
variable "aws_region" {
  type        = string
  description = "AWS region to deploy EKS cluster"
  default     = "us-east-1"
  validation {
    condition     = contains(["us-east-1", "us-west-2", "eu-west-1"], var.aws_region)
    error_message = "EKS 1.32 with Graviton4 is only supported in us-east-1, us-west-2, eu-west-1 as of Q3 2024."
  }
}

variable "cluster_name" {
  type        = string
  description = "Name of the EKS cluster"
  default     = "kubeflow-graviton4-cluster"
}

variable "cluster_version" {
  type        = string
  description = "Kubernetes version for EKS cluster"
  default     = "1.32"
  validation {
    condition     = var.cluster_version == "1.32"
    error_message = "Kubeflow 1.8 requires Kubernetes 1.32 for native ARM64 CRD support."
  }
}

# Configure AWS provider with default tags for cost tracking
provider "aws" {
  region = var.aws_region
  default_tags {
    tags = {
      Project     = "kubeflow-graviton4"
      Environment = "production"
      ManagedBy   = "terraform"
    }
  }
}

# EKS cluster resource with 1.32 version and ARM64 node support
resource "aws_eks_cluster" "kubeflow" {
  name     = var.cluster_name
  role_arn = aws_iam_role.eks_cluster.arn
  version  = var.cluster_version

  vpc_config {
    subnet_ids = aws_subnet.private[*].id
    endpoint_private_access = true
    endpoint_public_access  = false # Restrict public access for production
  }

  # Enable logging for audit trails
  enabled_cluster_log_types = ["api", "audit", "authenticator", "controllerManager", "scheduler"]

  depends_on = [
    aws_iam_role_policy_attachment.eks_cluster_policy
  ]
}

# IAM role for EKS cluster
resource "aws_iam_role" "eks_cluster" {
  name = "${var.cluster_name}-cluster-role"

  assume_role_policy = jsonencode({
    Version = "2012-10-17"
    Statement = [
      {
        Action = "sts:AssumeRole"
        Effect = "Allow"
        Principal = {
          Service = "eks.amazonaws.com"
        }
      }
    ]
  })
}

resource "aws_iam_role_policy_attachment" "eks_cluster_policy" {
  role       = aws_iam_role.eks_cluster.name
  policy_arn = "arn:aws:iam::aws:policy/AmazonEKSClusterPolicy"
}

# Managed node group with Graviton4 instances (c8g, m8g are Graviton4 family)
resource "aws_eks_node_group" "graviton4_nodes" {
  cluster_name    = aws_eks_cluster.kubeflow.name
  node_group_name = "graviton4-workload-nodes"
  node_role_arn   = aws_iam_role.eks_node.arn
  subnet_ids      = aws_subnet.private[*].id

  # Graviton4 instance types: c8g (compute optimized), m8g (general purpose)
  instance_types = ["c8g.2xlarge", "m8g.4xlarge"]

  scaling_config {
    desired_size = 2
    max_size     = 10
    min_size     = 1
  }

  # Use ARM64 AMI for Graviton compatibility
  ami_type       = "AL2023_ARM_64_STANDARD"
  capacity_type  = "ON_DEMAND" # Switch to SPOT for 70% cost savings in dev
  disk_size      = 100

  # Labels to target Graviton nodes for Kubeflow workloads
  labels = {
    "node.kubernetes.io/instance-family" = "graviton4"
    "workload"                           = "kubeflow-training"
  }

  depends_on = [
    aws_iam_role_policy_attachment.eks_node_policy
  ]
}

# IAM role for EKS nodes
resource "aws_iam_role" "eks_node" {
  name = "${var.cluster_name}-node-role"

  assume_role_policy = jsonencode({
    Version = "2012-10-17"
    Statement = [
      {
        Action = "sts:AssumeRole"
        Effect = "Allow"
        Principal = {
          Service = "ec2.amazonaws.com"
        }
      }
    ]
  })
}

# Attach required policies to node role
resource "aws_iam_role_policy_attachment" "eks_node_policy" {
  for_each = toset([
    "arn:aws:iam::aws:policy/AmazonEKSWorkerNodePolicy",
    "arn:aws:iam::aws:policy/AmazonEKS_CNI_Policy",
    "arn:aws:iam::aws:policy/AmazonEC2ContainerRegistryReadOnly",
    "arn:aws:iam::aws:policy/AmazonS3ReadOnlyAccess" # For training data access
  ])
  role       = aws_iam_role.eks_node.name
  policy_arn = each.value
}

# Output kubeconfig for cluster access
output "kubeconfig" {
  value = templatefile("${path.module}/kubeconfig.tpl", {
    cluster_name = aws_eks_cluster.kubeflow.name
    endpoint     = aws_eks_cluster.kubeflow.endpoint
    ca_crt       = aws_eks_cluster.kubeflow.certificate_authority[0].data
  })
  sensitive = true
}

#!/bin/bash
# Install Kubeflow 1.8 on Kubernetes 1.32 with Graviton4 support
# Exit on any command failure, treat unset variables as errors
set -euo pipefail
IFS=$'\n\t'

# Configuration variables - modify these for your environment
KUBEFLOW_VERSION="1.8.0"
KUBEFLOW_NAMESPACE="kubeflow"
AWS_REGION="us-east-1"
CLUSTER_NAME="kubeflow-graviton4-cluster"
GRAVITON_NODE_LABEL="node.kubernetes.io/instance-family=graviton4"

# Check for required dependencies
check_dependencies() {
  local deps=("kubectl" "aws" "git" "kustomize")
  for dep in "${deps[@]}"; do
    if ! command -v "$dep" &> /dev/null; then
      echo "ERROR: $dep is not installed. Please install it before proceeding."
      exit 1
    fi
  done
  echo "All dependencies satisfied."
}

# Validate kubectl connectivity to EKS cluster
validate_cluster_access() {
  echo "Validating cluster access..."
  if ! kubectl cluster-info &> /dev/null; then
    echo "ERROR: Cannot connect to EKS cluster. Run 'aws eks update-kubeconfig --name $CLUSTER_NAME --region $AWS_REGION'"
    exit 1
  fi
  # Check Kubernetes version is 1.32
  local k8s_version
  k8s_version=$(kubectl version --short | grep Server | awk '{print $3}' | cut -d. -f1,2)
  if [ "$k8s_version" != "v1.32" ]; then
    echo "ERROR: Kubernetes version $k8s_version detected. Kubeflow 1.8 requires v1.32."
    exit 1
  fi
  # Check for Graviton4 nodes
  if ! kubectl get nodes -l "$GRAVITON_NODE_LABEL" | grep -q "Ready"; then
    echo "ERROR: No Graviton4 nodes found with label $GRAVITON_NODE_LABEL. Check node group configuration."
    exit 1
  fi
  echo "Cluster access validated. Kubernetes version: $k8s_version, Graviton4 nodes found."
}

# Install Kubeflow 1.8 using official kustomize manifests
install_kubeflow() {
  echo "Downloading Kubeflow $KUBEFLOW_VERSION manifests..."
  local manifest_dir="kubeflow-manifests-$KUBEFLOW_VERSION"
  if [ -d "$manifest_dir" ]; then
    rm -rf "$manifest_dir"
  fi
  git clone --branch "v$KUBEFLOW_VERSION" https://github.com/kubeflow/manifests.git "$manifest_dir"
  cd "$manifest_dir" || exit 1

  echo "Applying Kubeflow CRDs (requires Kubernetes 1.32 for ARM64 compatibility)..."
  kustomize build example/kustomize/cluster-scoped-resources | kubectl apply -f -
  # Wait for CRDs to be established
  kubectl wait --for condition=established --timeout=60s crd -l app.kubernetes.io/part-of=kubeflow

  echo "Applying Kubeflow namespace-scoped resources..."
  kustomize build example/kustomize/namespaced-resources | kubectl apply -f -
  # Patch all Kubeflow deployments to run on Graviton4 nodes
  echo "Patching Kubeflow deployments to target Graviton4 nodes..."
  kubectl patch deployments -n "$KUBEFLOW_NAMESPACE" --all -p "{\"spec\":{\"template\":{\"spec\":{\"nodeSelector\":{\"$GRAVITON_NODE_LABEL\": \"\"}}}}}"
  # Wait for all deployments to roll out
  kubectl rollout status deployments -n "$KUBEFLOW_NAMESPACE" --all --timeout=300s
  cd ..
  echo "Kubeflow $KUBEFLOW_VERSION installed successfully in $KUBEFLOW_NAMESPACE namespace."
}

# Install Kubeflow Pipelines (KFPs) with ARM64 container support
install_kfp() {
  echo "Installing Kubeflow Pipelines $KUBEFLOW_VERSION..."
  local kfp_manifest="https://github.com/kubeflow/pipelines/releases/download/$KUBEFLOW_VERSION/install-manifest.yaml"
  kubectl apply -n "$KUBEFLOW_NAMESPACE" -f "$kfp_manifest"
  # Wait for KFP API server to be ready
  kubectl wait --for=condition=ready pod -l app=ml-pipeline -n "$KUBEFLOW_NAMESPACE" --timeout=300s
  echo "Kubeflow Pipelines installed. Access UI via kubectl port-forward -n $KUBEFLOW_NAMESPACE svc/ml-pipeline-ui 8080:80"
}

# Main execution flow
main() {
  echo "Starting Kubeflow $KUBEFLOW_VERSION installation on EKS 1.32 with Graviton4"
  check_dependencies
  validate_cluster_access
  install_kubeflow
  install_kfp
  echo "Full installation complete. Next step: Deploy training pipeline (see Section 4)."
}

main

# PyTorch training pipeline for Kubeflow 1.8 on Kubernetes 1.32 with Graviton4
# Requires kfp>=2.0.0 (Kubeflow 1.8 compatible) and torch>=2.3.0 (ARM64 support)
import kfp
from kfp import dsl
from kfp.dsl import Dataset, Input, Model, Output, component
import logging
from typing import Optional

# Configure logging for pipeline debugging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Define component to download CIFAR-10 training data to S3
@component(
    base_image="python:3.11-slim-bookworm-arm64", # ARM64 base image for Graviton4
    packages_to_install=["boto3>=1.34.0", "torchvision>=0.18.0"]
)
def download_cifar10(
    s3_bucket: str,
    s3_prefix: str,
    aws_region: str = "us-east-1"
) -> Output[Dataset]:
    """Download CIFAR-10 dataset and upload to S3 for training access."""
    import boto3
    from torchvision import datasets, transforms
    import os

    try:
        logger.info(f"Downloading CIFAR-10 dataset to {s3_bucket}/{s3_prefix}")
        # Define CIFAR-10 transform
        transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
        ])
        # Download training set
        train_set = datasets.CIFAR10(
            root="/tmp/cifar10",
            train=True,
            download=True,
            transform=transform
        )
        # Upload to S3 using boto3
        s3_client = boto3.client("s3", region_name=aws_region)
        # Save dataset to local path then upload (simplified for example)
        local_path = "/tmp/cifar10/cifar10-train.tar.gz"
        # In production, use torch.save or optimized serialization
        import tarfile
        with tarfile.open(local_path, "w:gz") as tar:
            tar.add("/tmp/cifar10/cifar10/train", arcname="train")
        s3_client.upload_file(local_path, s3_bucket, f"{s3_prefix}/cifar10-train.tar.gz")
        logger.info(f"Uploaded CIFAR-10 to s3://{s3_bucket}/{s3_prefix}/cifar10-train.tar.gz")
        return local_path
    except Exception as e:
        logger.error(f"Failed to download/upload CIFAR-10: {str(e)}")
        raise RuntimeError(f"Data preparation failed: {str(e)}") from e

# Define PyTorch training component for Graviton4 nodes
@component(
    base_image="pytorch/pytorch:2.3.0-cpu-arm64", # Official ARM64 PyTorch image
    packages_to_install=["boto3>=1.34.0", "torchvision>=0.18.0", "kfp>=2.0.0"]
)
def train_pytorch_model(
    input_dataset: Input[Dataset],
    s3_bucket: str,
    s3_prefix: str,
    epochs: int = 10,
    learning_rate: float = 0.001,
    batch_size: int = 64
) -> Output[Model]:
    """Train a simple CNN on CIFAR-10 using PyTorch on Graviton4."""
    import boto3
    import torch
    import torch.nn as nn
    import torch.optim as optim
    from torchvision import datasets, transforms
    import os
    import tarfile
    from kfp import dsl

    try:
        logger.info(f"Starting PyTorch training for {epochs} epochs on Graviton4 ARM64")
        # Verify ARM64 architecture (sanity check for Graviton4)
        import platform
        if platform.machine() != "aarch64":
            raise RuntimeError(f"Expected aarch64 (ARM64) but got {platform.machine()}")
        # Download dataset from S3
        s3_client = boto3.client("s3", region_name="us-east-1")
        local_data_path = "/tmp/cifar10-train.tar.gz"
        s3_client.download_file(s3_bucket, f"{s3_prefix}/cifar10-train.tar.gz", local_data_path)
        # Extract dataset
        with tarfile.open(local_data_path, "r:gz") as tar:
            tar.extractall("/tmp/cifar10-extracted")
        # Define simple CNN model
        class CIFAR10CNN(nn.Module):
            def __init__(self):
                super().__init__()
                self.conv1 = nn.Conv2d(3, 16, 3, padding=1)
                self.conv2 = nn.Conv2d(16, 32, 3, padding=1)
                self.pool = nn.MaxPool2d(2, 2)
                self.fc1 = nn.Linear(32 * 8 * 8, 128)
                self.fc2 = nn.Linear(128, 10)
                self.relu = nn.ReLU()
            def forward(self, x):
                x = self.pool(self.relu(self.conv1(x)))
                x = self.pool(self.relu(self.conv2(x)))
                x = x.view(-1, 32 * 8 * 8)
                x = self.relu(self.fc1(x))
                x = self.fc2(x)
                return x
        # Initialize model, loss, optimizer
        model = CIFAR10CNN()
        criterion = nn.CrossEntropyLoss()
        optimizer = optim.Adam(model.parameters(), lr=learning_rate)
        # Load data (simplified for example)
        transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
        train_set = datasets.CIFAR10("/tmp/cifar10-extracted", train=True, download=False, transform=transform)
        train_loader = torch.utils.data.DataLoader(train_set, batch_size=batch_size, shuffle=True)
        # Training loop
        for epoch in range(epochs):
            running_loss = 0.0
            for inputs, labels in train_loader:
                optimizer.zero_grad()
                outputs = model(inputs)
                loss = criterion(outputs, labels)
                loss.backward()
                optimizer.step()
                running_loss += loss.item()
            logger.info(f"Epoch {epoch+1}/{epochs}, Loss: {running_loss/len(train_loader):.4f}")
        # Save model to output path
        model_path = "/tmp/cifar10-cnn.pth"
        torch.save(model.state_dict(), model_path)
        # Upload model to S3
        s3_client.upload_file(model_path, s3_bucket, f"{s3_prefix}/cifar10-cnn.pth")
        logger.info(f"Model saved to s3://{s3_bucket}/{s3_prefix}/cifar10-cnn.pth")
        return model_path
    except Exception as e:
        logger.error(f"Training failed: {str(e)}")
        raise RuntimeError(f"PyTorch training failed: {str(e)}") from e

# Define the full Kubeflow pipeline
@dsl.pipeline(
    name="cifar10-pytorch-graviton4-pipeline",
    description="Train CIFAR-10 CNN on Graviton4 using Kubeflow 1.8 on K8s 1.32"
)
def cifar10_training_pipeline(
    s3_bucket: str = "kubeflow-training-data",
    s3_prefix: str = "cifar10",
    epochs: int = 10,
    learning_rate: float = 0.001,
    batch_size: int = 64
):
    # Step 1: Download CIFAR-10 data
    download_task = download_cifar10(
        s3_bucket=s3_bucket,
        s3_prefix=s3_prefix
    )
    # Target Graviton4 nodes for download task (lightweight, but consistent)
    download_task.set_node_selector({"node.kubernetes.io/instance-family": "graviton4"})

    # Step 2: Train PyTorch model
    train_task = train_pytorch_model(
        input_dataset=download_task.output,
        s3_bucket=s3_bucket,
        s3_prefix=s3_prefix,
        epochs=epochs,
        learning_rate=learning_rate,
        batch_size=batch_size
    )
    # Target Graviton4 nodes for training (compute intensive)
    train_task.set_node_selector({"node.kubernetes.io/instance-family": "graviton4"})
    # Mount S3 volume for faster data access (optional, for large datasets)
    train_task.add_env_variable("AWS_REGION", "us-east-1")

# Compile the pipeline to YAML for deployment
if __name__ == "__main__":
    try:
        logger.info("Compiling Kubeflow pipeline to cifar10-pipeline.yaml")
        kfp.compiler.Compiler().compile(
            pipeline_func=cifar10_training_pipeline,
            package_path="cifar10-pipeline.yaml"
        )
        logger.info("Pipeline compiled successfully. Upload to Kubeflow or run via kfp.Client()")
    except Exception as e:
        logger.error(f"Pipeline compilation failed: {str(e)}")
        raise

# Karpenter Provisioner for Graviton4 instances
apiVersion: karpenter.sh/v1beta1
kind: Provisioner
metadata:
  name: graviton4-training
spec:
  requirements:
    - key: karpenter.sh/capacity-type
      operator: In
      values: ["on-demand", "spot"]
    - key: node.kubernetes.io/instance-family
      operator: In
      values: ["c8g", "m8g", "r8g"] # Graviton4 families
    - key: kubernetes.io/arch
      operator: In
      values: ["arm64"]
  limits:
    resources:
      cpu: "1000"
      memory: "4000Gi"
  provider:
    subnetSelector:
      karpenter.sh/discovery: kubeflow-graviton4-cluster
    securityGroupSelector:
      karpenter.sh/discovery: kubeflow-graviton4-cluster
    tags:
      Project: kubeflow-graviton4
  ttlSecondsAfterEmpty: 60 # Scale down idle nodes after 1 minute

# Kaniko component to build ARM64 training container
@component(
    base_image="gcr.io/kaniko-project/executor:arm64-v1.20.0"
)
def build_training_image(
    dockerfile_path: str,
    context_path: str,
    ecr_repo: str,
    image_tag: str = "latest"
):
    import subprocess
    import os
    # Run Kaniko build for ARM64
    cmd = [
        "/kaniko/executor",
        "--dockerfile", dockerfile_path,
        "--context", context_path,
        "--destination", f"{ecr_repo}:{image_tag}-arm64",
        "--force",
        "--cleanup"
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        raise RuntimeError(f"Kaniko build failed: {result.stderr}")
    return f"{ecr_repo}:{image_tag}-arm64"

# Enable S3 caching for a Kubeflow component
@component(
    base_image="pytorch/pytorch:2.3.0-cpu-arm64",
    packages_to_install=["boto3>=1.34.0"]
)
def train_with_cache(
    input_data: Input[Dataset],
    learning_rate: float,
    epochs: int
) -> Output[Model]:
    # Define cache key based on input data hash and hyperparameters
    import hashlib
    cache_key = hashlib.sha256(
        f"{input_data.path}{learning_rate}{epochs}".encode()
    ).hexdigest()
    # Set KFP cache key (KFP 2.0+ supports this via annotation)
    import kfp
    kfp.dsl.get_pipeline_run_context().set_annotation(
        "pipelines.kubeflow.org/cache-key", cache_key
    )
    # Training logic here...

kubeflow-graviton4-tutorial/
├── terraform/               # EKS 1.32 + Graviton4 cluster config (Code Example 1)
│   ├── main.tf
│   ├── variables.tf
│   └── outputs.tf
├── scripts/                 # Installation and utility scripts (Code Example 2)
│   ├── install-kubeflow.sh
│   └── validate-cluster.sh
├── pipelines/               # Kubeflow pipeline code (Code Example 3)
│   ├── cifar10_pipeline.py
│   └── cifar10-pipeline.yaml
├── components/              # Reusable Kubeflow components
│   ├── download_cifar10/
│   └── train_pytorch/
├── kustomize/               # Kubeflow 1.8 kustomize patches
│   └── graviton4-node-selector.yaml
└── README.md                # Full tutorial steps and benchmarks