A Learning Method with Gap-Aware Generation for Heterogeneous DAG Scheduling

Abstract

Efficient scheduling of directed acyclic graphs (DAGs) in heterogeneous environments is challenging due to resource capacities and dependencies. In practice, the need for adaptability across environments with varying resource pools and task types, alongside rapid schedule generation, complicates these challenges. We propose WeCAN, an end-to-end reinforcement learning framework for heterogeneous DAG scheduling that addresses task--pool compatibility coefficients and generation-induced optimality gaps. It adopts a two-stage single-pass design: a single forward pass produces task--pool scores and global parameters, followed by a generation map that constructs schedules without repeated network calls. Its weighted cross-attention encoder models task--pool interactions gated by compatibility coefficients, and is size-agnostic to environment fluctuations. Moreover, widely used list-scheduling maps can incur generation-induced optimality gaps from restricted reachability. We introduce an order-space analysis that characterizes the reachable set of generation maps via feasible schedule orders, explains the mechanism behind generation-induced gaps, and yields sufficient conditions for gap elimination. Guided by these conditions, we design a skip-extended realization with an analytically parameterized decreasing skip rule, which enlarges the reachable order set while preserving single-pass efficiency. Experiments on computation graphs and real-world TPC-H DAGs demonstrate improved makespan over strong baselines, with inference time comparable to classical heuristics and faster than multi-round neural schedulers.

Figures (8)

• Figure 1: The illustration for spaces and maps. (a) The schedule space, divided into disjoint fibers for feasible schedule orders. (b) The serial schedule generation scheme $S_{\text{SGS}}$ embeds the feasible schedule order space $\mathcal{B}_f$ as a subset of the schedule space $\mathcal{X}$, which is "orthogonal" to each fiber $T^{-1}(w)$ as only intersect at one point $S_{\text{SGS}}(w)$. $S_{\text{SGS}}$ maps each schedule order to a best schedule (with minimal objective value) in the corresponding fiber. This relationship fix the optimality gaps. (c) The list scheduling map $S_{\text{list}}$ maps each schedule order to a best schedule in some fiber, but the fiber may not be the corresponding fiber. Therefore, the fiber containing the optimal solution may be missed since the optimal schedule order may be mapped to the schedule with other schedule orders. This mismatch leads to an optimality gap. (d) The skip action map $S_d$ is defined on a larger space $\overline{\Omega}_{\text{action}}\supset \Omega_{\text{action}}$, where the extended part $S_d(\overline{\Omega}_{\text{action}}\backslash \Omega_{\text{action}})$ covers the missed part $S_{\text{SGS}}(\mathcal{B}_f)\backslash S_{\text{list}}(\Omega_{\text{action}})$, making $TS_d$ a surjection to $\mathcal{B}_f$. The extended part closes the optimality gap.
• Figure 2: A counterexample showing list scheduling excludes the optimal solution for the DAG scheduling problem with heavy tasks. (a) Problem instance $\mathcal{P}_0$ with only one pool offering one type of resource with capacity 3. The resource demand and processing time is shown and the compatibility coefficients $K\equiv 1$. Each arrow $v\to w$ corresponds to an edge $(v,w)\in E$. (b) The only solution $x_0$ in the image of list scheduling $S_{\text{list}}(\Omega_{\text{action}}) = \{x_0\}$. The horizontal axis represents the time while the vertical axis represents the resource usage. (c) Optimal solution $x^*$ of the problem. (d) Worse solution for unrestricted waiting, leading to a significantly large objective value (makespan).
• Figure 3: The two-stage framework of our architecture. The first stage involves once network processing by our efficient and adaptive weighted cross-attention architecture. It embeds the problem attribution and provides scores and coefficients for the second stage. The second stage employs a generation map extended from list scheduling. This extension restores surjectivity onto the schedule order space and ensures the inclusion of an optimal solution in the image of the generation map, without significantly slowing down training convergence.
• Figure 4: Ablation study results of the skip action design. (a-b) Results of different skip action score formulations on TPC-H-30 and TPC-H-50 datasets with heavy tasks, where the improvement rate and the error bars of standard variation are labeled. "NoSkip" corresponds to the variant without skip action. "ConstSkip" corresponds to the variant with constant-score skip action. "DecrSkip" corresponds to results with our decreasing-score skip action. (c) Relative improvement of enabling the skip action in single-pass settings across scenarios with varying percentages of heavy tasks.
• Figure 5: The extension by introducing skip action with a constant score does not include optimal solution in its image for $\mathcal{P}_0$. It can not consistently close the optimality gap.

Theorems & Definitions (8)

• proof
• proof
• proof
• proof
• proof
• proof
• proof
• proof