Grid-Interactive Thermal Management of AI Data Centers via Contextual Distributionally Robust Optimization
For operators of AI data centers, this work addresses the challenge of balancing thermal safety and grid flexibility under uncertain workloads, offering a practical improvement over existing robust control methods.
This paper proposes a Contextual Distributionally Robust Optimization (CDRO) framework for thermal management in AI data centers, which dynamically adapts uncertainty bounds using real-time context. The method achieves near-zero thermal violations under extreme workload spikes and reduces the operational cost premium of robustness by approximately 13.7 percentage points compared to standard Min-Max MPC.
Thermal management in AI data centers is increasingly challenged by bursty workloads and uncertain heat generation. To prevent thermal violations, existing cooling strategies either enforce conservative, rigid bounds that severely limit grid responsiveness, or rely on forecast-driven controllers that perform poorly under AI workload uncertainty and distribution shifts. To overcome the above challenges, this paper proposes a Contextual Distributionally Robust Optimization (CDRO) framework for grid-interactive cooling control. Unlike standard DRO with fixed ambiguity sets, the proposed approach dynamically adapts the Wasserstein radius using real-time AI and grid context. This safely shrinks uncertainty bounds during stable regimes, unlocking deep demand-side flexibility. Theoretically, we formulate the control as an infinite-dimensional inf-sup problem, derive an exact tractable reformulation for the Wasserstein worst-case expected-cost term, and then derive a tractable conservative deterministic counterpart for the Distributionally Robust Conditional Value at Risk (DR-CVaR) thermal safety constraint. Solved via a scalable nested Alternating Direction Method of Multipliers (ADMM) algorithm, the CDRO controller achieves near-zero thermal violations under extreme workload spikes in high-fidelity EnergyPlus co-simulations. Simultaneously, it reduces the operational cost premium of robustness by approximately 13.7 percentage points relative to standard Min-Max Model Predictive Control (MPC).