Title: GAA-PtrNet: Graph attention aggregation-based pointer network for one-shot DAG scheduling
Primary area: optimization
Keywords: DAG  Scheduling, Graph Attention, Pointer Network, Reinforcement Learning, Combinatorial Optimization

Abstract: Optimizing Directed Acyclic Graph (DAG) workflow makespan by scheduling techniques is a critical issue in the high performance computing area. Many studies in recent years combined Pointer Network (PtrNet) with reinforcement learning (RL) to schedule DAGs by generating DAG task priorities in a sequence-to-sequence manner. However, these PtrNet-based scheduling methods need to repeatedly compute the decoder's hidden state or context embeddings according to the recent local decisions, which leads to limited capability of exploiting the DAG global topological structure, high computation complexity and inability to achieve one-shot scheduling. To address these issues, we propose GAA-PtrNet, a novel PtrNet based on graph attention aggregation (GAA) for one-shot DAG workflow scheduling. In GAA-PtrNet, we compute the pair-wise graph attention scores among nodes in one-shot, then directly aggregate these scores to obtain the probability of selecting candidate nodes. Consequently, the explicit decoder or context embedding structure in PtrNet is omitted in our GAA-PtrNet, and the network takes only one shot forward propagation to infer a solution for a whole DAG scheduling problem, significantly reducing the computation complexity. Additionally, to train GAA-PtrNet, we design a training strategy based on policy gradient RL with dense reward signal and demonstration learning. To our knowledge, GAA-PtrNet is the first network model to achieve PtrNet-based one-shot DAG scheduling. GAA-PtrNet can better handle with DAG workflow structures, providing high quality DAG scheduling solutions. The experimental results show that the proposed method is superior in terms of objective and runs about 10 times faster when compared to previous PtrNet-based methods, and also performs better than other learning-based DAG scheduling methods.