Reinforcement Learning for Heterogeneous DAG Scheduling with Weighted Cross-Attention

Ruisong Zhou, Haijun Zou, Li Zhou, Chumin Sun, Zaiwen Wen

Keywords: Combinatorial Optimization, Directed Acyclic Graph, Heterogeneous Scheduling, Reinforcement Learning

Abstract: Efficient scheduling of directed acyclic graphs (DAGs) in heterogeneous environments is challenging due to diverse resource capacities and intricate dependencies. In practice, the need for adaptability across environments with varying resource pools, task types, and other settings, alongside rapid schedule generation, complicates these challenges. We propose WeCAN, an end-to-end reinforcement learning framework for heterogeneous DAG scheduling featuring task-resource compatibility. WeCAN rapidly generates schedules through single-pass network inference. Leveraging the weighted cross-attention layer, WeCAN utilizes all available environment information while preserving adaptability across diverse heterogeneous environments. Moreover, we analyze the optimality gap inherent in list-scheduling-based methods, revealing their inability to guarantee optimal solutions and their reduced performance in certain cases. Under the single-pass setting, we develop a method to enable skip actions, addressing this gap without sacrificing computational efficiency. Our approach delivers robust performance and adaptability, outperforming state-of-the-art methods across diverse datasets.

Primary Area: optimization

Submission Number: 4310