Skip to content
datarekha

Infrastructure as code for ML

Turn GPU pools, buckets, IAM, and endpoints from undocumented clicks into reviewable, repeatable infrastructure.

13 min read Intermediate MLOps Lesson 33 of 35

What you'll learn

  • Explain why declarative infrastructure and idempotent changes are safer than imperative cloud scripts
  • Read a plan diff, protect stateful resources, and understand remote state locking
  • Promote one ML platform module through dev, staging, and production with separate state
  • Provision GPU pools with spot capacity, labels, and taints without confusing node scheduling
  • Detect drift and keep secrets, model artifacts, and training data out of infrastructure state

Before you start

At 02:57 on a Tuesday morning, the prediction API is returning 503s.

The model worked in staging. Production has the same container image, Kubernetes manifest, and application configuration. Yet its pods cannot schedule: the GPU node pool has the wrong instance family and incorrect scheduling metadata, so the cluster has no suitable GPU capacity.

The other differences are separate failures. Public subnets can create routing or security problems. A workload IAM role that cannot read the model bucket causes a pod that starts to fail during model download. A missing node-bootstrap permission can stop nodes joining the cluster. Those are runtime and provisioning failures, not scheduler rejection.

Nobody changed the deployment that evening. Three people changed the cloud console over the past month.

That is the gap infrastructure as code, usually shortened to IaC, is meant to close. IaC puts intended infrastructure in reviewable files that can be tested and applied repeatedly. The files explain why production looks the way it does, instead of forcing someone to reconstruct it from browser history.

For an ML platform, this includes networks, buckets, IAM roles, clusters, GPU pools, queues, registries, and inference endpoints. Model code is only a passenger. IaC builds the road.

Desired state beats a recipe

An imperative script says what to do, in order:

  1. Create a bucket.
  2. Create a role.
  3. Create a cluster.
  4. Add nodes.
  5. Attach permissions.

If it stops halfway through, rerunning may find an existing bucket, a differently named role, or a partially created node group. The script knows its steps; it does not necessarily know the current world.

A declarative configuration says what the world should contain:

  • one versioned artifact bucket;
  • one IAM role with a defined policy;
  • one GPU node group, ranging from zero to four nodes;
  • a taint that keeps ordinary workloads away.

Terraform compares that desired state with real resources and chooses API calls to close the gap. You describe the destination; the tool calculates the route.

Idempotent means applying the same input twice produces the same result as applying it once. Setting bucket versioning to enabled is idempotent: whether it is already enabled or not, the final result is enabled. This matters because cloud calls time out, CI retries, and humans rerun commands.

Mathematically, an idempotent operation satisfies f(f(x)) = f(x). IaC does not guarantee idempotency: providers have bugs, APIs are eventually consistent, and some resources require replacement. It does provide a desired end state and a way to inspect proposed changes first.

The production loop is:

Git desiredconfigurationPlanReview plandiff and risksApply
The plan is a proposal. Apply changes the cloud only after someone accepts that proposal.

A plan refreshes the cloud view and shows a diff: + creates, ~ updates in place, - destroys, and -/+ replaces. A replacement means the provider cannot change a field in place.

A worked ML platform

Suppose a team trains a fraud model nightly and serves it during the day. It needs a versioned S3 artifact bucket, an existing EKS cluster, a g5.xlarge GPU node group using spot capacity, zero to four GPU nodes, a GPU taint, and a separate read-only workload role.

This small Terraform example assumes the EKS cluster, node IAM role, private subnets, and workload identity setup already exist. A usable GPU path also requires an EKS-compatible NVIDIA AMI and a working device plugin or operator.

terraform {
  required_version = ">= 1.5.0"

  required_providers {
    aws = {
      source = "hashicorp/aws"
    }
  }
}

provider "aws" {
  region = var.region
}

variable "region" {
  type    = string
  default = "us-east-1"
}

variable "cluster_name" {
  type = string
}

variable "node_role_arn" {
  type = string
}

variable "private_subnet_ids" {
  type = list(string)
}

variable "artifact_bucket_name" {
  type = string
}

resource "aws_s3_bucket" "artifacts" {
  bucket = var.artifact_bucket_name

  tags = {
    Purpose = "ml-artifacts"
  }

  lifecycle {
    prevent_destroy = true
  }
}

resource "aws_s3_bucket_versioning" "artifacts" {
  bucket = aws_s3_bucket.artifacts.id

  versioning_configuration {
    status = "Enabled"
  }
}

resource "aws_eks_node_group" "gpu_spot" {
  cluster_name    = var.cluster_name
  node_group_name = "${var.cluster_name}-gpu-spot"
  node_role_arn   = var.node_role_arn
  subnet_ids      = var.private_subnet_ids

  # Use the value appropriate for the cluster version and architecture.
  ami_type = "AL2023_x86_64_NVIDIA"

  capacity_type  = "SPOT"
  instance_types = ["g5.xlarge"]

  scaling_config {
    desired_size = 0
    min_size     = 0
    max_size     = 4
  }

  # Use this only when Cluster Autoscaler owns desired_size.
  lifecycle {
    ignore_changes = [scaling_config[0].desired_size]
  }

  labels = {
    accelerator = "nvidia-a10g"
    workload    = "training"
  }

  taint {
    key    = "workload"
    value  = "gpu"
    effect = "NO_SCHEDULE"
  }
}

output "artifact_bucket" {
  value = aws_s3_bucket.artifacts.bucket
}

output "gpu_node_group" {
  value = aws_eks_node_group.gpu_spot.node_group_name
}

Terraform gives the bucket and node group stable addresses, aws_s3_bucket.artifacts and aws_eks_node_group.gpu_spot, and maps them to real cloud IDs in state.

The NVIDIA AMI matters: an ordinary EKS image can boot a GPU instance without drivers, leaving the node with no usable GPU resource. The device plugin or GPU operator then discovers the hardware and advertises nvidia.com/gpu to Kubernetes. Until both pieces work, a pod requesting a GPU has nowhere to run.

A first plan might say:

Plan: 3 to add, 0 to change, 0 to destroy.

That means this configuration adds one bucket, one versioning resource, and one node group; it is not a performance benchmark.

Changing instance_types to another family may require node-group replacement. Replacement can drain workloads and interrupt training. A bucket replacement is more dangerous.

prevent_destroy = true makes an accidental Terraform destroy fail loudly. It does not protect against direct console deletion or make data indestructible.

The scaling values are not autoscaling by themselves. Cluster Autoscaler must discover this node group and raise its desired count when pending GPU pods cannot fit, up to four, then reduce it when idle. In that design, ignore_changes prevents Terraform from fighting the autoscaler over the live desired count. If Terraform owns desired count, remove the rule and keep other controllers away. Karpenter is an alternative ownership model; do not casually configure both.

The node role lets nodes join EKS and perform node operations. Workloads need a separate role through EKS Pod Identity or IRSA, with least-privilege access such as s3:GetObject for the model prefix. Add s3:ListBucket only if the loader lists objects. If the binding is wrong, expect AccessDenied or 403 during model download—not a scheduling error.

State is the tool’s memory

Terraform state maps configuration addresses to real resource identities and selected attributes. It may contain IDs, endpoints, network details, and policy documents, so treat it as sensitive operational data.

A local state file is unsuitable for team production work. Use a remote backend with encryption, strict permissions, version history or recovery, audit logging, and locking.

If Alice and Ben both read state version 42 and apply simultaneously without locking, each can write a different version 43. The cloud may contain both changes while state records only one. The next plan then proposes changes that already happened.

With locking, one apply proceeds and the other waits or fails with Error acquiring the state lock. Check whether the first apply is active. Clear a stale lock only after confirming no process is using the state. Force-unlock is a scalpel, not a retry button.

Dev and production need separate state, distinct credentials, and environment-scoped permissions. Separate backend keys isolate records but are not a security boundary: broad credentials can still reach another environment. Separate accounts or projects provide stronger boundaries when the operational cost is justified.

Modules make promotion boring

A module is reusable infrastructure configuration. Put common platform structure in one module and give each environment different inputs:

infra/
  modules/
    ml-platform/
      main.tf
      variables.tf
      outputs.tf
  envs/
    dev/
      main.tf
      terraform.tfvars
    staging/
      main.tf
      terraform.tfvars
    prod/
      main.tf
      terraform.tfvars

For example, dev might allow zero to four GPU nodes, staging zero to eight, and production two to 32 across multiple zones with stricter IAM. Those values encode cost, availability, and blast-radius policy.

Promote a reviewed Git commit from dev to staging to production. The module owns common structure; environment roots own deliberate differences. Expose real decisions—region, subnets, instance family, scaling bounds, tags—not every provider detail.

What belongs in IaC

IaC should own durable infrastructure and permissions: networks, buckets and lifecycle rules, IAM, clusters and node pools, registries, queues, databases, endpoints, and observability integrations.

It should usually not own their contents. The bucket belongs in IaC; an 18-gigabyte model artifact and a 4.2-terabyte training dataset do not. Those need separate versioning, lineage, retention, and promotion workflows. Terraform can provision the registry and the identity that reads it.

Likewise, IaC can create an EKS cluster and node pools, while Kubernetes manifests or Helm manage Deployments, Services, and Jobs. Application images and model releases belong in a deployment workflow such as GitHub Actions for ML, not in a cloud-infrastructure apply. See data and model versioning and Kubernetes for ML for those boundaries.

GPU-specific scheduling

Spot capacity may be unavailable in a requested zone or family even when the plan is valid. Spot nodes can also disappear; training needs checkpoints and retries. Serving generally belongs on stable capacity unless interruption is tolerable.

Labels and taints solve different problems:

  • A label is metadata selected by a pod.
  • A taint repels pods unless they have a matching toleration.
  • A selector or affinity chooses nodes, while a GPU request reserves the device.

A training pod therefore needs a toleration, a GPU label selector, and nvidia.com/gpu: 1:

apiVersion: batch/v1
kind: Job
metadata:
  name: fraud-trainer
spec:
  template:
    spec:
      serviceAccountName: trainer
      restartPolicy: Never
      tolerations:
        - key: workload
          operator: Equal
          value: gpu
          effect: NoSchedule
      nodeSelector:
        accelerator: nvidia-a10g
      containers:
        - name: trainer
          image: example.com/fraud-trainer:2026-08-28
          resources:
            limits:
              nvidia.com/gpu: 1

The device plugin must advertise the resource before scheduling can succeed.

If an autoscaler changes desired count, assign ownership explicitly. Otherwise Terraform sees the useful change as drift and continually undoes it. Avoid broad ignore_changes, which hides real failures.

Drift and secrets

Drift is a difference between declared and real infrastructure. A normal plan refreshes provider state and may show, for example, a console change from GPU maximum four to eight being restored to four. Detection is not repair: CI can alert without applying. Decide whether the console change should be encoded in Git or reverted.

Provider defaults, asynchronous service changes, and intentionally ignored fields can complicate drift. Constrain provider versions and review lock-file changes.

Do not put passwords, access keys, tokens, or signing keys in Terraform arguments. Providers may write values to state. sensitive = true hides output but does not guarantee absence from state.

Create secret containers and access policies with IaC, but write secret values through a controlled secrets-management process. Let workloads fetch them at runtime, and use short-lived CI identity such as OIDC federation. See ML security.

A secret name or ARN can be infrastructure. The secret value is runtime data.

Common first symptoms

  • -/+ for a bucket, database, or endpoint: stop and inspect the immutable field, resource address, and state identity. Consider a state move or import only after confirming the real object.
  • State-lock error: find the active apply or verify a stale lock before clearing it.
  • GPU pods pending: inspect scheduler events, taints, tolerations, labels, GPU requests, device-plugin logs, instance capacity, and node-group maximums.
  • Every plan reverses an autoscaler change: two systems own one field; choose one owner.
  • A token appears in plan or state: rotate it, remove payload management from IaC, and inspect state history and CI logs.

The honest limitation

IaC gives repeatability, not correctness. A reviewed configuration can still request unavailable GPU capacity, grant excessive IAM permissions, exceed quota, or create a network that cannot reach the model registry. A plan does not prove that the application starts, a model fits in memory, or spot interruption is survivable.

The abstraction also costs state administration, provider lag, careful migrations, and code review. That is a good trade for shared infrastructure serving thousands of requests or consuming expensive GPU time, but a poor one for a disposable twelve-minute experiment with no shared dependencies.

What to remember

  • Declarative IaC describes an end state; idempotency makes retries converge safely.
  • A plan is a proposed diff. Investigate every destroy and replacement.
  • Remote state needs encryption, access control, versioning, and locking; environment isolation also requires distinct credentials.
  • Keep infrastructure in IaC, but keep artifacts, datasets, and secret values in their own systems.
  • GPU scheduling requires capacity planning, drivers, device discovery, labels, taints, tolerations, and a GPU request.

Quick check

0/3
Q1
Q2
Q3

Sign in to track your progress

Completed lessons, your XP, level, and streak save to your account — it's free and takes a few seconds.

Practice this in an interview

All questions
How do you attribute and control ML spend across teams and models (FinOps for ML)?

Apply FinOps to ML by tagging every workload (training jobs, endpoints, GPU pools) by team, model, and environment so cost is attributable, then track unit-economics metrics like cost per prediction or per training run rather than just total spend. Set budgets and alerts, identify idle GPUs and overprovisioned endpoints, and enforce guardrails like autoscaling and instance-type policies. The goal is continuous visibility and accountability so teams optimize cost without killing experimentation.

How does CI/CD for ML differ from standard software CI/CD, and what stages should an ML pipeline include?

ML CI/CD must validate not just code correctness but also model quality — automated retraining triggers, data validation, model evaluation gates, and canary deployment checks that standard software pipelines have no equivalent for. A regression in model AUC is as much a deployment failure as a 500 error.

How does autoscaling work for ML inference services, and what metrics should drive it?

Autoscale ML inference on a leading workload signal such as per-pod queue wait or backlog, with GPU or model-specific throughput as supporting signals; CPU alone misses accelerator saturation. Kubernetes HPA supports custom or external metrics, while strict online SLAs usually require warm replicas and scale-to-zero only for workloads that tolerate measured cold starts.

Why use a pipeline orchestrator like Airflow or Kubeflow instead of cron scripts for ML workflows?

ML workflows are multi-step DAGs with dependencies, and an orchestrator gives you dependency management, retries, backfills, caching, observability, and lineage that chained cron jobs cannot. Airflow is a general-purpose task orchestrator defining DAGs in Python, while Kubeflow Pipelines is ML-native, passing typed artifacts between containerized steps on Kubernetes with conditional logic like deploy only if accuracy exceeds a threshold. Choosing depends on whether you need generic scheduling or ML-specific, container-based pipelines.

Related lessons

Explore further