Generalizable Graph Reinforcement Learning for Task Scheduling under Workload Distribution Shifts

Authors

  • Ruohan Xie School of Computer Science, Nanjing University, Nanjing, Jiangsu, China
  • Yiming Cao School of Computer Science, Nanjing University, Nanjing, Jiangsu, China

Keywords:

task scheduling, distributed systems, graph reinforcement learning, workflow DAGs, distribution shift, heterogeneous clusters

Abstract

Learning-based task schedulers promise adaptive decisions for distributed workloads represented as directed acyclic graphs (DAGs). Their practical value, however, depends on whether a scheduler trained under one workload mix remains effective when task sizes, runtime tails, resource demands, and cluster heterogeneity change. This paper presents GRACE-DS, a graph reinforcement scheduler that combines DAG-criticality features, cluster context, and distribution-randomized policy search to improve zero-shot generalization under workload distribution shifts. Instead of learning a policy only for a fixed trace, GRACE-DS trains a residual graph policy over task-node candidates across randomized workload domains and uses normalized structural features to preserve scale transfer. A discrete-event simulator evaluates five operating regimes: normal workloads, larger DAGs, heavy-tailed runtimes, resource-demand shifts, and heterogeneous clusters. Across three independent seeds, GRACE-DS reduces mean makespan by 0.89% to 2.45% relative to the strongest baseline in every shift regime, while maintaining competitive utilization. The results show that modest graph-structured reinforcement can complement strong list scheduling heuristics and provide a reproducible foundation for robust scheduling research under distribution shift.

References

Yue, X., & Zhu, R. (2024). An explainable hybrid machine learning framework for cash-flow-conditioned multi-horizon liquidity risk early warning: An SME-oriented benchmark study. World Journal of Management Science, 2(1), 63-72.

Jiao, Y., & Liang, Y. (2024). An integrated process mining and explainable machine learning framework for bottleneck and internal control risk detection in the multi entity financial close. Social Science and Management, 1(3), 96-107.

Bian, M., Wang, Y., & Lin, Y. (2024). Explainable supplier risk prediction for critical telecommunications infrastructure procurement: An empirical XGBoost–SHAP framework. Social Science and Management, 1(1), 64-77.

Liang, Y., & Jiao, Y. (2024). Real-time financial fraud detection using a feature-light temporal dynamic graph neural network with contrastive learning. Journal of Trends in Finance and Economics, 1(1), 76–84.

Shang, W., & Chen, T. (2024). PrefixKD-Edge: Real-time locomotion-intent prediction with prefix self-distillation on lightweight temporal networks. Innovation and Technology Studies, 1(1), 30–38.

Rhee, M., Zou, J., & Qin, Y. (2024). CLARITY: A closed loop framework for duplicate aware failure prediction and budgeted automated remediation in cloud infrastructure. AI and Data Science Journal, 1(1), 76-85.

J. Dean and S. Ghemawat, "MapReduce: Simplified data processing on large clusters," Communications of the ACM, vol. 51, no. 1, pp. 107-113, 2008.

A. Tumanov, J. Cipar, G. R. Ganger, and M. A. Kozuch, "TetriSched: Global rescheduling with adaptive plan-ahead in dynamic heterogeneous clusters," in Proc. USENIX ATC, 2016, pp. 241-254.

R. Grandl, M. Chowdhury, A. Akella, and G. Ananthanarayanan, "GRAPHENE: Packing and dependency-aware scheduling for data-parallel clusters," in Proc. OSDI, 2016, pp. 81-97.

V. K. Vavilapalli et al., "Apache Hadoop YARN: Yet another resource negotiator," in Proc. SoCC, 2013, pp. 1-16.

B. Hindman et al., "Mesos: A platform for fine-grained resource sharing in the data center," in Proc. NSDI, 2011, pp. 295-308.

K. Ousterhout, P. Wendell, M. Zaharia, and I. Stoica, "Sparrow: Distributed, low latency scheduling," in Proc. SOSP, 2013, pp. 69-84.

A. Mirhoseini et al., "Device placement optimization with reinforcement learning," in Proc. ICML, 2017, pp. 2430-2439.

R. S. Sutton and A. G. Barto, Reinforcement Learning: An Introduction, 2nd ed. Cambridge, MA, USA: MIT Press, 2018.

T. N. Kipf and M. Welling, "Semi-supervised classification with graph convolutional networks," in Proc. ICLR, 2017.

W. L. Hamilton, R. Ying, and J. Leskovec, "Inductive representation learning on large graphs," in Proc. NeurIPS, 2017, pp. 1024-1034.

P. Velickovic et al., "Graph attention networks," in Proc. ICLR, 2018.

J. Gilmer et al., "Neural message passing for quantum chemistry," in Proc. ICML, 2017, pp. 1263-1272.

S. Yang, G. Ding, Z. Chen, and J. S. Yang, "GART: Graph neural network-based adaptive and robust task scheduler for heterogeneous distributed computing," IEEE Access, vol. 13, pp. 200196-200216, 2025, doi: 10.1109/ACCESS.2025.3633290.

Z. Wu et al., "A comprehensive survey on graph neural networks," IEEE Transactions on Neural Networks and Learning Systems, vol. 32, no. 1, pp. 4-24, 2021.

T. Coleman, H. Casanova, L. Pottier, M. Kaushik, E. Deelman, and R. Ferreira da Silva, "WfCommons: A framework for enabling scientific workflow research and development," Future Generation Computer Systems, vol. 128, pp. 16-27, 2022.

S. Bharathi et al., "Characterization of scientific workflows," in Proc. WORKS, 2008, pp. 1-10.

Qin, Y., Zou, J., & Rhee, M. (2022). An AI-driven intelligent transaction execution framework for internet finance platforms: Cost-sensitive real-time integrity interdiction. Eurasia Journal of Science and Technology, 4(2), 36–44.

Wang, Y., & Zhang, C. (2023). AERIS: Adaptive executor scoped resource management for privacy preserving data processing in Apache Spark. Eurasia Journal of Science and Technology, 5(2), 25-33.

Fang, Q. (2024). True-EWS: Trustworthy early warning with equalized conformal guarantees for equitable student support in U.S. education. Educational Research and Human Development, 1(2), 44–54.

Shen, A., & Huang, S. (2022). A systems-aware framework for intelligent large-scale recommendation models. Journal of Computer Science and Electrical Engineering, 4(1), 36–44.

Downloads

Published

2026-09-03