Infrastructure as Code for AI: Terraform & Pulumi for ML Platforms

Why IaC for AI Infrastructure?

AI infrastructure is complex and critical: training clusters with GPUs, serving endpoints with auto-scaling, feature stores, model registries, experiment tracking servers. Without IaC: environments drift apart (works on staging, fails on prod), infrastructure knowledge lives in one person's head, environment recreation takes days, no audit trail for changes.

With IaC: environments are identical and reproducible, infrastructure changes go through code review, recreate any environment in minutes, full audit trail in Git history.

Terraform for ML Infrastructure

AWS SageMaker with Terraform

Provision a SageMaker training job infrastructure: SageMaker execution role with S3 and ECR permissions, VPC and subnets for training (optional, for sensitive data), training job with specific instance type (ml.p3.8xlarge for GPU training), model artifact destination in S3, and endpoint configuration for serving.

Terraform resource structure: aws_iam_role for SageMaker execution, aws_s3_bucket for data and artifacts, aws_sagemaker_model with primary container (ECR image), aws_sagemaker_endpoint_configuration with production variants and auto-scaling, aws_sagemaker_endpoint to deploy.

GPU Training Cluster

Terraform for EC2 GPU training: aws_placement_group with cluster strategy for low-latency NVLink, aws_launch_template with p3.8xlarge, 100GB gp3 EBS, user data to install CUDA/PyTorch, aws_autoscaling_group for on-demand GPU instances (scale to 0 when not training).

Spot instances for training cost reduction: request spot instances via aws_spot_instance_request or launch template. Use EC2 Fleet with multiple instance types for best availability. Implement checkpointing to resume from interruptions.

Kubeflow on EKS

EKS cluster for ML workloads: aws_eks_cluster with managed node groups (CPU for serving, GPU for training), aws_eks_node_group with g5 instances for GPU workloads, aws_iam_role for IRSA (IAM Roles for Service Accounts). Install Kubeflow via Terraform Helm provider.

Pulumi for ML Infrastructure

Pulumi advantages over Terraform: real programming languages (Python, TypeScript), loops and conditionals without HCL workarounds, testing infrastructure code with pytest/jest, component packages for reusable infrastructure.

Python example: create a Kubernetes cluster resource, add GPU node pool, deploy MLflow Helm chart with S3 backend for artifact storage. Component class: MLPlatform creates all required resources and exposes endpoints as outputs.

Multi-Environment Strategy

Environment Matrix

Dev: small instances (no GPU), shared infrastructure, aggressive cost optimization. Staging: mirrors production configuration, real data samples, performance testing. Production: high availability, auto-scaling, multi-region failover, monitoring.

Terraform workspaces or separate directories for each environment. Variables file per environment: instance_type, min_nodes, max_nodes, multi_az, backup_retention.

Module Structure

Modules for reusable components: module "ml_training_cluster" (provisions GPU ASG, spot fleet, IAM roles), module "model_serving" (EKS deployment, HPA, ALB ingress, Route53). Call modules from environment-specific main.tf with environment-specific variables.

GitOps for ML Infrastructure

Terraform Cloud / Atlantis

All Terraform changes via pull requests. Atlantis runs plan automatically on PR, posts plan output as PR comment. Require approval before apply. Apply automatically on merge to main. Full audit trail: who changed what, when, and why.

Drift Detection

Scheduled Terraform plan in CI (daily): detect when infrastructure drifts from code. Alert Slack when drift detected. Require correction PR before next deployment.

Cost Management for AI Infrastructure

GPU instances are expensive. Cost controls: CloudWatch + Lambda auto-shutdown for idle training clusters, Spot instances for non-critical training (70-90% cost reduction), Reserved Instances/Savings Plans for sustained serving load.

Terraform tags: tag all resources with environment, team, cost_center, project. Use AWS Cost Explorer with tag-based filtering for detailed ML cost attribution.

Infrastructure as Code transforms ML infrastructure from art to engineering—reproducible, reviewable, and reliable.