When you're training a model across hundreds of GPUs, things break in weird ways. Node failures, network partitions, NCCL timeout cascades, OOM conditions, silent gradient corruption. ClusterOrch-Gym is a reinforcement learning environment where agents learn to detect these failures and execute recovery policies — autonomously, at machine speed, without waiting for a human operator to page in.
The Environment
Discrete-Event Simulator
Models real cluster telemetry from actual distributed training runs. The simulator ingests raw NCCL logs, dmesg outputs, GPU metrics from nvidia-smi, and network telemetry from infiniband counters. Event timestamps are drawn from empirical distributions so the failure arrival patterns statistically match production clusters.
State Space
GPU utilization per device, memory pressure (allocated vs. reserved vs. free), network latency between node pairs, NCCL collective timing (all-reduce, all-gather, reduce-scatter completion times), and checkpoint freshness (seconds since last successful checkpoint). The full observation vector is 847-dimensional for a 256-GPU cluster.
Action Space
Checkpoint rollback to last known good state, topology reconfiguration (ring vs. tree vs. hierarchical all-reduce), power capping of thermal-throttled nodes, job migration to a reserved spare pool, and node isolation for suspected hardware faults. The action space is structured as a hierarchical policy: high-level actions select a recovery strategy, low-level actions parameterize it (e.g., which checkpoint to roll back to).
Reward Function
Minimize training downtime while maximizing resource utilization. The reward is a weighted combination: negative time spent in degraded state, positive tokens/second processed, negative penalty for unnecessary preemption (false positives), and a sparsity bonus for correctly diagnosing intermittent faults on the first attempt. Weights are tuned via inverse reinforcement learning from operator logs.
Training
Agents are trained using PPO (Proximal Policy Optimization) with a Transformer-based policy network that processes the temporal structure of telemetry data. The policy network uses a causal attention mask so it can only condition on past observations, mirroring the real-time constraints of the deployment environment. Training runs on a separate GPU cluster (4x A100s) with 32 parallel environment instances, achieving ~2M environment steps per hour.
Agents are evaluated on failure scenarios extracted from real-world distributed training logs across three cluster sizes: 8 GPUs (single-node), 64 GPUs (multi-node), and 256 GPUs (cluster scale).
Failure Modes
The benchmark includes 15 distinct failure modes organized into 4 categories:
Hardware — GPU XID errors, PCIe link degradation, ECC memory errors exceeding threshold, fan failure causing thermal throttling.
Network — NCCL watchdog timeout, infiniband link flap, TCP incast congestion, ARP table exhaustion on top-of-rack switches.
Software — CUDA out-of-memory, NCCL version mismatch across nodes, deadlocked all-reduce due to topology mismatch, checkpoint I/O hang on NFS.
Silent — Gradient corruption (NaN injection with no error signal), straggler node slowing collective ops without failing, silent data corruption from ECC-disabled memory.