A Priority-Driven Deep Reinforcement Learning Scheduler for Resource Management

Deep neural networks (DNNs) have gained popularity in various fields, including computer vision and speech recognition. However, the growing size of data and complexity of models have made training DNNs increasingly time-consuming. To accelerate computation, parallel and distributed training of DNNs using multiple GPUs is widely adopted. This approach, however, introduces significant challenges in managing and scheduling GPU resources. One critical challenge is job prioritization, which ensures that urgent jobs are processed promptly. Current solutions often overlook prioritization settings, leading to suboptimal performance. For example, neglecting prioritization can result in long delays for urgent jobs. Conversely, overly prioritizing urgent jobs without balancing the needs of other jobs can lead to extended completion times (JCT) for less urgent tasks. While existing approaches that incorporate prioritization typically rely on heuristic-based methods, which may not be optimal, deep reinforcement learning (DRL)-based methods offer greater adaptability and can generally outperform traditional heuristic-based approaches. To address these challenges, we introduce Prioritize, a novel GPU job scheduler designed to enhance GPU resource management in distributed DNN training. Prioritize employs a DRL model to efficiently select jobs based on their priority. To minimize the JCT for lower-priority jobs, Prioritize dynamically increases their priority over time, thereby reducing their queueing delays. The DRL model maximizes a priority-based reward to ensure effective job selection. Extensive experiments demonstrate that Prioritize outperforms several baseline methods in terms of queueing time, JCT, and Makespan. The proposed approach not only achieves superior scheduling performance but also effectively addresses prioritization challenges, making Prioritize a promising solution for efficient GPU job scheduling in distributed DNN training scenarios.