[Paper Review] Shockwave: Fair and Efficient Cluster Scheduling for Dynamic Adaptation in Machine Learning

Pengfei Zheng, Rui Pan+3 more|arXiv (Cornell University)|Sep 30, 2022

Advanced Bandit Algorithms Research4 citations

TL;DR

Shockwave is a proactive, market-theory-based scheduler that co-optimizes fairness and efficiency in GPU clusters running machine learning jobs with dynamic resource adaptation. By extending Fisher market equilibrium to dynamic, time-varying utility functions and using stochastic dynamic programming with Bayesian prediction, Shockwave reduces makespan by 1.3× and improves fairness by 2× over existing schedulers in trace-driven and real-cluster experiments.

ABSTRACT

Dynamic adaptation has become an essential technique in accelerating distributed machine learning (ML) training. Recent studies have shown that dynamically adjusting model structure (e.g., lottery ticket hypothesis) or hyperparameters (e.g., batch size) can significantly accelerate training without sacrificing accuracy. However, existing ML cluster schedulers are not designed to handle dynamic adaptation. We show that existing schemes fail to provide fairness and degrade system efficiency when the training throughput changes over time under dynamic adaptation. We design Shockwave, a scheduler with future planning that builds on two key ideas. First, Shockwave extends classic market theory from static settings to dynamic settings to co-optimize efficiency and fairness. Second, Shockwave utilizes stochastic dynamic programming to handle dynamic changes. We build a system for Shockwave and validate its performance with both trace-driven simulation and cluster experiments. Results show that for traces of ML jobs with dynamic adaptation, Shockwave improves makespan by 1.3X and fairness by 2X when compared with existing fair scheduling schemes.

Motivation & Objective

• To address the lack of fairness and efficiency in existing cluster schedulers when handling ML jobs with dynamic adaptation (e.g., changing batch sizes).
• To design a scheduler that proactively plans for future resource demands rather than reacting to them.
• To co-optimize long-term fairness and system efficiency in the presence of time-varying job utility and resource requirements.
• To ensure robustness against prediction errors and dynamic changes in job behavior.
• To validate the scheduler’s superiority through trace-driven simulation and real cluster experiments.

Proposed method

• Formulates a discrete-time dynamic Fisher market where each job has equal budget and the market clears to achieve equilibrium, ensuring efficiency and fairness.
• Extends classic market theory to handle time-varying job utilities (e.g., throughput changes due to batch size scaling) using stochastic dynamic programming.
• Employs Bayesian statistics and interpolation to predict job runtimes and schedule slacks under dynamic adaptation.
• Uses a timeout-controlled solver (e.g., Gurobi) to compute near-optimal allocations with bounded gap from optimality, minimizing scheduling overhead.
• Introduces a Nash social welfare objective that penalizes skewed job progress, promoting conservative, fair allocation decisions.
• Proactively schedules jobs across rounds by planning for future utility changes, avoiding reactive re-prioritization that harms efficiency.

Experimental results

Research questions

• RQ1 Can a market-based scheduler maintain fairness and efficiency when ML jobs dynamically change their resource demands over time?
• RQ2 How does proactive scheduling with future planning compare to reactive or static scheduling in dynamic adaptation workloads?
• RQ3 To what extent does prediction error in job runtime estimation affect the fairness and efficiency of the scheduler?
• RQ4 Can a dynamic market formulation with time-varying utilities achieve better system-wide performance than existing fair schedulers like DRF, Themis, or AlloX?
• RQ5 How does solver overhead and real-time performance impact the practical deployment of such a scheduler in production clusters?

Key findings

• Shockwave reduces makespan by 1.3× compared to state-of-the-art fair schedulers like Themis, Gavel, and AlloX in trace-driven simulations.
• Shockwave improves fairness (measured by FTF and fraction of unfairly scheduled jobs) by 2× over existing schedulers under dynamic adaptation.
• The solver overhead is minimal—bound gap is only 0.03% at 15s timeout for 500 jobs and 0.44% for 2000 jobs—showing high scalability and low latency.
• Shockwave demonstrates resilience to prediction errors: fairness degrades slowly even with ±100% noise in runtime estimates, maintaining performance comparable to baselines.
• The scheduler’s proactive planning significantly outperforms reactive approaches, as reactive re-prioritization fails to compensate for early inefficiencies.
• Shockwave’s use of Nash social welfare ensures strong fairness by penalizing imbalanced job progress, leading to conservative and equitable scheduling decisions.