Large-Scale GPU Infrastructure for Neural Web Search Training

exa.ai · Tech · 2025 · source Agentic: no · promotional bias: low · source-backed · extracted with prompt 0.5.6 (schema v2)

Generated from extraction.json — do not edit by hand. Every claim cites verbatim source evidence; stated means the quote directly asserts the claim, implied means it supports it indirectly.

Problem

Domain: GPU infrastructure for neural search model training

Exa needed a repeatable GPU infrastructure stack that could turn bare-metal H200 and A100 clusters into reliable training and batch-processing capacity for neural web retrieval. The platform had to support embedding billions of pages, reranker training, checkpoint caching with S3 backups, concurrent ML workflows, multi-node PyTorch jobs, cloud bursting, developer debugging, priority scheduling, lineage, and reproducible environments without manual GPU juggling.

Constraints:

• the new cluster arrived as 18 bare-metal GPU nodes
• the combined fleet had 224 GPUs across A100 and H200 clusters
• ML engineers were running many training experiments in parallel
• inference requirements had grown 100x and the ML team had grown 5x
• training jobs needed local iteration, multi-node execution, cloud bursting, and reproducibility
• training data needed fast on-prem access while keeping S3 as the durable source of truth
• hardware and driver management needed to avoid manual SSH and node-specific drift

Architecture

Exa built Hephaestus, a five-layer cluster stack around Pulumi, Ansible/Kubespray, NVIDIA GPU and Network Operators, Alluxio, and Flyte. Pulumi provides typed infrastructure-as-code and orchestrates bare-metal and cloud resources. Ansible/Kubespray prepares servers and joins them into Kubernetes. NVIDIA operators install drivers, CUDA, container runtimes, device plugins, GPUDirect RDMA, RoCE, and metrics. Alluxio turns the fleet's NVMe into a distributed cache with S3 as the authority. Flyte schedules code-first ML workflows, multi-node PyTorch/MPI jobs, Ray/KubeRay batch jobs, cloud burst queues, checkpoint/resume, task caching, lineage, priority queues, developer pods, log routing, metrics, and artifact storage.

Diagram confidence: reconstructed (reconstructed from prose).

flowchart LR
  a100["A100 cluster"] --> k8s["Kubernetes fleet"]
  h200["H200 Exacluster"] --> k8s
  pulumi["Pulumi infrastructure code"] --> ansible["Ansible Kubespray automation"]
  ansible --> k8s
  k8s --> nvidia["NVIDIA operators"]
  k8s --> flyte["Flyte workflow platform"]
  k8s --> alluxio["Alluxio cache"]
  alluxio --> s3["S3 durable store"]
  flyte --> alluxio
  flyte --> s3
  flyte --> cloud["Cloud burst queues"]
  nvidia --> observability["Prometheus and Loki observability"]
  flyte --> observability
  pulumi --> cloud

Components

Data flow

Pulumi provisions infrastructure and runs Ansible/Kubespray to transform powered-on bare-metal servers into Kubernetes nodes. NVIDIA operators install and maintain the GPU/network runtime layer. Training data and checkpoints are addressed through S3 URLs; Alluxio serves hot objects from distributed NVMe and flushes or fetches to S3 as needed. Flyte registers Docker-image, requirements, and source-code-hashed tasks, then schedules workflows onto on-prem or cloud queues with topology-aware multi-node placement. Logs route to Loki, metrics to Prometheus, artifacts to S3 with Alluxio caching, and lineage links inputs, outputs, and exact code versions.

Why this architecture

Exa built a Kubernetes-centered, code-first stack because the team needed one repeatable path from bare metal to production ML workflows, not ad-hoc scripts or manual GPU scheduling. Pulumi and Kubespray make provisioning auditable and reproducible; NVIDIA operators prevent node drift; Alluxio keeps hot training data local while retaining S3 durability; and Flyte matched the full workflow requirement list, including multi-node training, cloud burst, checkpoint/resume, lineage, and priority queues.

Models used: retrieval models · reranker · 7B parameter pre-train · PyTorch DDP · PyTorchJob · MPIJob · Ray · KubeRay · Flyte

Scale

• Exa purchased a $5 million GPU cluster.
  “Exa recently purchased a 5 million dollar GPU cluster”article, stated
• The new H200 Exacluster has 144 H200 GPUs.
  “144xNVIDIA H200-141 GB (8 per node)”article, stated
• The combined fleet has 224 GPUs.
  “224 GPUs”article, stated
• The combined fleet has 26.4 TB of GPU RAM.

  "Total GPU RAM

6.4 TB

20 TB

26.4 TB"

— article, stated

• The combined fleet has about 350 TB of local NVMe.

  "Local NVMe

80 TB

270 TB

350 TB"

— article, stated

• A run that used to take a week can now finish in about one day.
  “a run that used to take a week to complete can now finish in ~1 day.”article, stated
• The cluster is used to embed billions of pages.
  “we're simultaneously embedding billions of pages, training a reranker, uploading checkpoints to a cache with S3 backups”article, stated
• The stack can launch a 7B parameter pre-train on 16 nodes.
  “Launch a 7B parameter pre-train on 16 nodes with a single command.”article, stated
• The platform can queue 30 hyperparameter sweeps in parallel.
  “Queue 30 hyper-parameter sweeps in parallel—no manual GPU juggling.”article, stated
• The whole cluster can be rebuilt to 224 ready GPUs in under one hour.
  “Rebuild the entire cluster—from bare metal to 224 GPUs ready—in under one hour.”article, stated

Alternatives considered (roads not taken)

Manage cluster setup with ad-hoc scripts

Rejected because: Exa needed a repeatable, auditable process for 28 bare-metal nodes, high-speed networks, and core services.

Manually manage GPU and network drivers on each node

Rejected because: NVIDIA operators provide uniform driver, CUDA, runtime, GPUDirect, RoCE, and metrics management without manual SSH or node-specific scripts.

Serve all training data directly from S3

Rejected because: Alluxio caches hot data on local NVMe to avoid upload/download bottlenecks and egress costs while preserving S3 as source of truth.

Use Kubeflow, Airflow, Prefect, Dagster, or similar schedulers

Rejected because: Exa says Flyte was the only evaluated framework covering their full workflow list without extensive custom work.

Patterns

durable-execution (workflow)

Flyte provides checkpointing, partial-run resume, task caching, and re-runnable versioned workflow registrations for long-running ML jobs.

batch-offline-processing (training)

The platform runs large embedding, reranking, hyperparameter sweep, pre-training, and cloud-burst batch jobs across GPU fleets.

semantic-caching (storage)

Alluxio acts as a distributed cache for training objects while S3 remains the durable source of truth.

observability-tracing (operations)

The platform routes logs and metrics while NVIDIA operators publish DCGM metrics for GPU health and Flyte tracks lineage.

cost-aware-routing (capacity)

Flyte can route the same project to on-prem or cloud queues, and Alluxio keeps hot data on-prem to reduce egress.

tool-use-orchestration (platform)

The stack composes Pulumi, Ansible/Kubespray, Kubernetes, NVIDIA operators, Alluxio, Flyte, Ray/KubeRay, PyTorchJob, MPIJob, Loki, Prometheus, and S3 into a unified ML platform.

Tradeoffs

Infrastructure-as-code format

Chose Pulumi with Python over YAML-only or ad-hoc scripts.

Rationale: Python gives engineers familiar control flow and tests, while Pulumi provides typed validation, diffs, and rollback.

Data cache authority

Chose Alluxio local NVMe cache with S3 as source of truth over Refactor code away from S3 URLs or use only local ephemeral storage.

Rationale: The team gets local-disk-speed hot data, cost control, transparent S3-compatible access, and durable S3 backing.

Workflow orchestrator

Chose Flyte over Kubeflow, Airflow, Prefect, Dagster, or custom orchestration.

Rationale: Flyte met the broad set of code-first, local, distributed, batch, cloud-burst, checkpoint, lineage, queue, and developer-debug requirements without extensive custom work.

Capacity placement

Chose On-prem GPU clusters with optional cloud burst over Only local execution or only cloud execution.

Rationale: The same project can use on-prem queues for owned capacity or scale to thousands of spot GPUs with a configuration change.

Gotchas (what bit them)

Bare-metal GPU fleets drift if setup is handled manually.

Root cause: Many nodes, high-speed networks, and core services create too much state for ad-hoc scripts and SSH sessions.

Resolution: Use Pulumi with Ansible/Kubespray and NVIDIA operators so provisioning, driver management, and upgrades are repeatable.

Training data throughput can bottleneck expensive GPUs.

Root cause: Direct object-store access and upload speed are too slow for hot training data at aggregate local-disk speeds.

Resolution: Use Alluxio to cache objects on distributed NVMe while keeping S3 durable and authoritative.

A GPU scheduler has to support research iteration and production-scale jobs at the same time.

Root cause: ML engineers need local debugging, multi-node training, cloud burst, lineage, checkpoints, and fair-share policies in one workflow system.

Resolution: Adopt Flyte because it covered the full requirement list without extensive custom work.

Hardware/software upgrades can disrupt training if rolled out fleet-wide.

Root cause: CUDA and driver changes affect GPU workloads and can create incompatibility with network operators.

Resolution: Use Kubernetes operators with version-aligned upgrades, cordon and drain one node at a time, then test before advancing.

Evaluation

Exa evaluates the platform through reproducible infrastructure diffs, type checking with mypy, Kubernetes node-health verification, driver/operator rollout testing, Flyte traceability and lineage, Prometheus/DCGM metrics, and concrete throughput outcomes such as pre-training launches, parallel sweeps, cloud-burst jobs, and full-cluster rebuild time.

Metrics: mypy validation of Pulumi parameters · node health verification after Kubernetes join · DCGM metrics · Prometheus metrics · Loki logs · dataset and artifact lineage · 7B pre-train on 16 nodes · 30 hyperparameter sweeps in parallel · 100-node AWS batch job · under-one-hour rebuild to 224 ready GPUs

Outcomes

• Exa built a combined 224-GPU training fleet across 80 A100s and 144 H200s.
• Workloads that used to take roughly a week can now finish in about one day.
• The cluster runs nearly 24/7 and concurrently supports embedding billions of pages, reranker training, checkpoint uploads, and other ML workflows.
• The platform can launch a 7B pre-train on 16 nodes, queue 30 hyperparameter sweeps, run 100-node AWS batch jobs, and rebuild the cluster to 224 ready GPUs in under an hour.
• Flyte gives engineers code-first workflows, task cache, partial-run resume, traceability, lineage, priority queues, and local debugging.

The infrastructure includes a $5M cluster, 28 nodes, 224 total GPUs, 26.4 TB GPU RAM, 46 TB system RAM, 350 TB local NVMe, roughly 168 FP16 PFLOPs, 100 KW operating power, billion-page embedding workflows, 7B parameter pre-training on 16 nodes, 30 concurrent sweeps, and 100-node cloud-burst jobs.