Generalizable Graph Reinforcement Learning for Task Scheduling under Workload Distribution Shifts
Keywords:
task scheduling, distributed systems, graph reinforcement learning, workflow DAGs, distribution shift, heterogeneous clustersAbstract
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
Issue
Section
License
Copyright (c) 2026 Software Engineering and Artificial Intelligence Systems

This work is licensed under a Creative Commons Attribution 4.0 International License.