NEURAL DAG SCHEDULING VIA ONE-SHOT PRIORITY SAMPLING

Abstract

We consider the problem of scheduling operations/nodes, the dependency among which is characterized by a Directed Acyclic Graph (DAG). Due to its NP-hard nature, heuristic algorithms were traditionally used to acquire reasonably good solutions, and more recent works have proposed Machine Learning (ML) heuristics that can generalize to unseen graphs and outperform the non-ML heuristics. However, it is computationally costly to generate solutions using existing ML schedulers since they adopt the episodic reinforcement learning framework that necessitates multi-round neural network processing. We propose a novel ML scheduler that uses a one-shot neural network encoder to sample node priorities which are converted by list scheduling to the final schedules. Since the one-shot encoder can efficiently sample the priorities in parallel, our algorithm runs significantly faster than existing ML baselines and has comparable run time with the fast traditional heuristics. We empirically show that our algorithm generates better schedules than both non-neural and neural baselines across various real-world and synthetic scheduling tasks.

1. INTRODUCTION

The problem of scheduling operations arises across many domains, such as data centers where the incoming jobs have to be scheduled on a distributed server (Mao et al., 2019) , manufacturing pipelines in the form of job shop scheduling problems (JSSP) (Manne, 1960) , and ML compilers where the operations of a computation graph need to be scheduled on the available hardware devices (Paliwal et al., 2020; Zhou et al., 2020) . In all these cases, the problem may be abstracted using a directed acyclic graph (DAG) where the nodes of the graph represent the operations and the edges represent the dependency constraints between the operations and hence the problem is also referred to as DAG scheduling. The objective is to minimize the finish time (or makespan) of the DAG subject to resource and dependency constraints. It is well known that this is an NP-hard problem (Kan, 2012), and practitioners have traditionally relied on heuristic methods to obtain good solutions. One of the celebrated scheduling approaches is list scheduling (Graham, 1969) where the idea is to schedule nodes as early as possible and to break ties using priorities. The priorities can be obtained via different node metrics which are computationally inexpensive such as critical-path based, shortest processing time or most operations remaining (Haupt, 1989) . More recently, researchers have proposed deep reinforcement learning based methods to solve scheduling problems (Zhang et al., 2020; Zhou et al., 2020; Wang et al., 2021; Mao et al., 2019) . The scheduling policy in all the references utilize Graph Neural Networks (GNN) as an encoder to derive node embeddings. Zhang et al. (2020) proposed an auto-regressive GNN based policy for the JSSP problem which predicts the next node for scheduling given the nodes scheduled so far. Wang et al. (2021) proposed a bi-level optimization approach which modifies the input DAG by adding multiple edges via a learned policy and then apply the critical-path heuristic on the modified DAG. One major drawback of the existing ML based schedulers is the computational cost as they require multi-round neural network processing (encoding step). The multi-round neural network processing is reflected as auto-regressive architecture (Zhang et al., 2020) or bi-level optimization

