A Predict-then-Schedule framework for Power Distribution Networks with AI Data Centers
2026-07-20 • Computational Engineering, Finance, and Science
Computational Engineering, Finance, and Science
AI summaryⓘ
The authors address the problem of high energy use and costs in AI data centers by improving how workload predictions are used to schedule tasks. They point out that just making better predictions doesn't always lead to better scheduling results. To fix this, they created a combined system called Predict-Then-Schedule (PTS) that learns to predict and schedule tasks together for better efficiency. Their system also considers power limits and penalties for shifting too much work at once. Tests show their method lowers costs and makes the system more reliable than older methods.
GPU workloadsAI data centersworkload schedulingpower grid conditionspredictive optimizationdifferentiable convex optimizationgradient-based trainingload-sheddingoperational cost
Authors
Siqi Yan, Jiebao Zhang, Xi Yao, Juan Huang, Ye Shi
Abstract
The surge of GPU-intensive workloads in artificial intelligence (AI) data centers drives massive energy demands, leading to soaring costs and significant stress on local power distribution networks. Coordinating delay-tolerant workload scheduling with power grid conditions via precise workload prediction can mitigate these issues. However, a critical gap remains in conventional approaches, i.e., minimizing prediction error does not necessarily lead to minimized downstream operational loss. Hence, this paper proposes an end-to-end Predict-Then-Schedule (PTS) framework that integrates upstream workload prediction with downstream scheduling optimization. By leveraging differentiable convex optimization, the PTS framework maps input features directly to optimal scheduling and enables gradient-based training. Furthermore, to respect the data center's capacity, a workload over-shifted loss combining electricity cost with a penalty for load-shedding is introduced to evaluate scheduling quality. Experiments demonstrate that the proposed framework significantly reduces operational cost and enhances system security compared to the conventional two-stage baseline.