PaperlayerView exact revision on arXiv

Source-linked AI summary

OAK: Restart- and Age-Aware Scheduling for Distributed Machine Learning on Shared GPU Clusters

Khaled Aljbab, Amine Barrak

arXiv:2609.19024v2cs.DC

TL;DR

Shared-cluster DML schedulers often optimize goodput or fairness without jointly treating accumulated waiting and restart cost as first-class signals. OAK adds age-aware and decomposed-restart terms to a per-round MILP; its decomposition outperforms an aggregate restart formulation, with the gap widening as failures increase.

  • Problem

    Existing goodput-driven schedulers do not make accumulated waiting a first-class signal, and Pollux aggregates productive time and checkpoint overhead instead of tracking them separately.

    Source: §I · p. 1

  • Method

    OAK adds a bounded age key for accumulated waiting and a decomposed restart factor for productive training and checkpoint overhead to Sia’s per-round MILP.

    Source: §I · p. 2, B. Composite Utility · p. 3, C. Restart Factor and Age Key · p. 4

  • Results

    OAK’s decomposed restart factor consistently outperforms the aggregate formulation, with the gap widening as failure rate rises; its age key compresses worst-case JCT on contention-heavy workloads.

    Source: D. Decomposed Restart Factor:… · p. 8, §VII · p. 11

  • Takeaways & Limitations

    Hybrid OAK shows that the composite utility can be used with reduced-population NSGA II while retaining genetic search’s high-failure reallocation robustness at a fraction of Pollux’s per-round decision cost.

    Source: §VII · p. 11

    Evaluation covers a 12-GPU simulated cluster and 4-GPU hardware validation, data-parallel jobs, and Poisson failures; production-scale behavior and other regimes remain untested.

    Source: §VI · p. 10

Abstract

from arXiv · show

Distributed machine learning increasingly runs on shared multi-tenant GPU clusters where contention and failures are routine. Goodput-driven schedulers maximise instantaneous throughput but treat accumulated waiting time and restart cost as second-class signals: long-waiting jobs are repeatedly deferred, and the minutes of checkpoint loading paid after each interruption are not folded back into allocation decisions. We present OAK, a per-round mixed-integer linear scheduler whose composite utility adds two first-class terms to goodput: an age key that elevates jobs with high cumulative waiting time, and a decomposed restart factor that tracks productive training time and checkpoint overhead as separately measured quantities rather than a single aggregate ratio. We evaluate OAK against four representative goodput- and fairness-driven baselines in a trace-driven simulator on a 12-GPU cluster across four data-parallel workloads with controlled Poisson failure injection, and reproduce the failure-mode improvement on a 4-V100 real-hardware cluster. Under failure-free conditions OAK matches the strongest goodput-driven baseline within 4%. Under failures at rate $λ= 0.1$ it reduces mean job completion time (JCT) by 33.9-58.3% and worst-case JCT by 58-76% over its goodput-driven foundation; the decomposed restart factor alone accounts for a 40-57% mean and 52-76% tail JCT improvement against an aggregate restart estimator used in prior goodput-driven schedulers. Per-round decision latency is 7.65 ms, three orders of magnitude faster than evolutionary-search alternatives at default settings.

I. INTRODUCTION

Shared-cluster DML schedulers face contention, long-running jobs, and routine failures, while existing goodput- and fairness-driven approaches do not combine goodput, accumulated waiting, and restart cost in one objective. OAK adds an age key and a decomposed restart factor to Sia’s per-round MILP, and evaluates the design in simulation and on real hardware.

I. INTRODUCTION

  • DML jobs can occupy dozens of GPUs for days or weeks, and failures and checkpoint-coupled preemption make restart time a material scheduling concern.Production measurements report contention, heterogeneity, and a non-trivial fraction of GPU-time consumed by failed or restarted jobs.Source: §I · p. 1

A. Cluster Schedulers for Deep-Learning Training

  • Goodput- and fairness-driven schedulers address parts of the problem, but none of the four baselines combines goodput, waiting time, and restart cost in one per-round objective.The related work groups cluster schedulers into allocation-efficiency and goodput- or fairness-driven approaches.Source: §I · p. 1, A. Cluster Schedulers for… · p. 2

I. INTRODUCTION

  • OAK adds an accumulated-waiting age key and a decomposed restart factor to Sia’s goodput-driven per-round MILP.The paper reports sub-10 ms decision time at 12-GPU scale.Source: §I · pp. 1–2
  • At λ = 0.1, the decomposed restart factor reduces mean JCT by 40–57% and worst-case JCT by 52–76% versus Pollux’s aggregate estimator.The decomposition separately tracks productive training time and checkpoint/restore overhead.Source: §I · p. 2
  • The evaluation compares OAK with Sia, Pollux, Gavel, and DRF in a 12-GPU trace-driven simulator across four data-parallel workloads, with 4-V100 hardware validation.The simulator comparison includes controlled Poisson failure injection, as described in the paper context.Source: §I · p. 2

B. Restart Cost and Recovery in DML Training

Runtime-level recovery systems reduce the cost of individual restart events, while OAK feeds accumulated checkpoint time into scheduling decisions. OAK extends Sia by changing its per-job, per-configuration score and retaining the rest of Sia’s design.

B. Restart Cost and Recovery in DML Training

  • Runtime and worker-level systems target individual checkpoint, recovery, or latency costs, but do not feed restart-induced inefficiency into the cluster allocation objective.OAK is complementary: its restart factor feeds accumulated checkpoint time into the per-round score rather than reducing the cost of a single restart.Source: B. Restart Cost and Recovery in DML… · pp. 2–3
  • OAK’s restart factor complements recovery systems by feeding accumulated checkpoint time back into the per-round score.This steers scheduling away from configurations that would generate further restart cost.Source: B. Restart Cost and Recovery in DML… · pp. 2–3

III. DESIGN OF OAK

  • OAK changes only the score that the ILP maximizes, adding age-key and restart-factor terms to Sia’s goodput term.It reuses Sia’s feasible configuration sets, trace-based goodput estimates, placement step, and optimization backend.Source: §III · p. 3

A. Scheduling Model and Notation

OAK schedules active jobs in rounds, using job-specific feasible configurations and temporal counters for waiting and checkpoint time. Its composite utility combines goodput, a restart factor, and an age key, then supports a linear integer optimization.

A. Scheduling Model and Notation

  • Each round begins on a fixed timer or an active-job-set change; GPUs are assigned to at most one job, using that job’s feasible configurations.The experiments use a 60-second timer.Source: A. Scheduling Model and Notation · p. 3
  • OAK tracks queue waiting time for the age key and checkpoint/restore time separately from productive training time for the restart factor.Productive training time is job age minus checkpoint/restore time.Source: A. Scheduling Model and Notation · p. 3

B. Composite Utility

  • The per-job/configuration score combines goodput and the restart factor under exponent p, then multiplies by an age key that boosts long-waiting jobs.The age key remains outside the power so its boost is independent of the goodput-shaping exponent.Source: B. Composite Utility · p. 3
  • The exponent p shapes the effect of goodput and restart efficiency; the paper uses p = 0.5 for soft fairness and restricts p to nonnegative values.The paper leaves independent shaping of goodput and restart efficiency to future work.Source: B. Composite Utility · p. 3
  • OAK selects binary job-configuration assignments, and precomputes the powered utility so the optimization remains a linear integer program.A job is left idle when all its feasible utilities fall below the allocation threshold.Source: B. Composite Utility · pp. 3–4

C. Restart Factor and Age Key

OAK combines a restart factor based on measured checkpoint overhead with an age key that boosts long-waiting jobs. The factor penalizes restart-heavy jobs while the age key can offset that penalty for jobs accumulating queue wait.

C. Restart Factor and Age Key

  • OAK’s restart factor uses productive training time, measured checkpoint time, and the per-application restart penalty; checkpoint recovery lowers the factor through both the numerator and denominator.For newly submitted jobs with Ai(t) < Prestart,i, OAK sets ri = 1 so they are not penalized before making progress.Source: C. Restart Factor and Age Key · p. 4
  • Unlike Pollux’s restart-count-times-penalty estimate, OAK uses measured checkpoint time, which remains exact at any failure rate.Pollux’s aggregate estimate can accumulate error as restarts increase.Source: C. Restart Factor and Age Key · p. 4
  • OAK’s capped exponential age key raises scores with queue waiting time, accelerating priority for chronically starved jobs while limiting any one job’s dominance.The experiments use α = 0.01 and Kmax = 100; the cap is described as a robustness guard above values reached in the workloads.Source: C. Restart Factor and Age Key · p. 4
  • In a three-job example, OAK admits the efficient J1 and chronically waiting J3, while Sia’s goodput-only choice is restart-heavy J2.With p = 0.5, OAK’s utilities are approximately 1.50, 1.11, and 64.6 for J1, J2, and J3, respectively.Source: C. Restart Factor and Age Key · p. 4

D. Per-Round Scheduling Workflow

OAK adds score generation and optimization steps to Sia’s round workflow. Because the score is separable from the optimizer, another optimizer accepting per-pair utility can replace the optimization steps without changing the score function.

D. Per-Round Scheduling Workflow

  • Steps 3–5 generate OAK’s score, while Steps 6–7 optimize it; any optimizer accepting per-pair utility can replace those optimization steps without changing the score function.Steps 1–2 and 6–8 are inherited unchanged from Sia; Steps 3–5 are OAK’s additions.Source: D. Per-Round Scheduling Workflow · p. 4

IV. METHODOLOGY

The methodology evaluates OAK and four baselines in a configurable simulator using profiled workloads, per-round scheduling, and controlled Poisson failures. It measures restart and queue-wait counters and uses those quantities in OAK’s scheduling factors.

A. Hardware and Simulator

  • The simulator runs on a three-node, 12-V100 cluster and compares schedulers under paired workloads, submission times, failure streams, and trace-driven goodput estimates.Each scheduler starts a round every 60 seconds; paired comparisons make relative JCT differences invariant to the simulator’s absolute throughput calibration.Source: A. Hardware and Simulator · pp. 4–5

B. Workloads

  • Workloads use profiled computer-vision, NLP, and speech training jobs across four contention scenarios ranging from 4 jobs to 40.The scenarios are workload-small, workload-medium, workload-congestion, and workload-stress.Source: B. Workloads · p. 5

C. Composite Utility Solver

  • At each round, OAK solves a MILP over binary allocation variables using precomputed utilities and feasible configurations filtered before optimization.At the stated cluster scale, solve times stay below 100 ms, within the 60-second scheduling interval.Source: C. Composite Utility Solver · p. 5

D. Failure Injection and Restart Factor

  • Poisson failures increment each job’s measured checkpoint-and-restore counter, while queue-wait time accrues during no allocation or non-progressing recovery.OAK substitutes measured checkpoint time for the restart-count estimate in its restart factor; newly submitted jobs receive ri = 1 when Ai < Prestart.Source: D. Failure Injection and Restart Factor · p. 5

E. Baselines

The evaluation compares OAK with goodput- and fairness-oriented schedulers and reports results across repeated random seeds. Seed counts vary by experiment, with larger samples for the principal failure-rate sweep.

E. Baselines

  • The four baselines span goodput optimization, multi-objective goodput and statistical efficiency, finish-time fairness, and dominant-resource fairness.They are Sia, Pollux, Gavel, and DRF, respectively.Source: E. Baselines · p. 5

F. Statistical Methodology

  • The main OAK–Sia comparison uses 10 seeds per cell, while the workload-medium and workload-congestion failure-rate sweep uses 30.Sensitivity sweeps and the four-scheduler comparison use 5 seeds per cell, except Pollux uses 10; schedulers in that comparison share seed values.Source: F. Statistical Methodology · pp. 5–6

V. OAK SCHEDULER EVALUATION

At λ = 0, OAK, Sia, and Pollux are within 4% of one another in mean JCT, while Gavel and DRF are slower. Calibration finds that the restart penalty is the sensitive parameter, unlike the age-key settings and fairness exponent.

1) Calibration Sweep Methodology:

  • The calibration independently sweeps p, α, Kmax, and Prestart on workload-medium at λ = 0.05, holding the other parameters fixed across five seeds per cell.Figure 2 summarizes the sweeps.Source: 1) Calibration Sweep Methodology · p. 6

2) Sensitivity Findings:

  • Mean JCT varies by less than 1.5% across the tested ranges for p, α, and Kmax; α saturates above 0.05 and Kmax is rarely reached.The chosen defaults are within 1% of each parameter’s per-sweep minimum.Source: 2) Sensitivity Findings · p. 6
  • Prestart changes mean JCT by 12% across the tested range, from 32,377 s at 10 s to 36,246 s at 120 s.Lower values allow more aggressive preemption when failures are rare, trading some failure resilience for throughput; this motivates per-application restart-cost lookup.Source: 2) Sensitivity Findings · p. 6
  • The experiments use p = 0.5, α = 0.01, and Kmax = 100, with a per-application Prestart table because one constant would over- or under-penalize applications with different recovery times.At α = 0.01, Ki reaches Kmax after roughly 460 seconds of accumulated waiting; p = 0 is excluded because it removes the goodput optimization gradient.Source: 3) Selected Defaults · p. 6, 2) Sensitivity Findings · p. 6

B. Performance under Failure-Free Conditions

  • Within 4%, OAK, Sia, and Pollux differ in mean JCT on every workload at λ = 0; Gavel is 1.7–6.7× slower and DRF is 1.2–6.8× slower than OAK.The comparison covers four workloads and five schedulers; Fig. 3 reports mean JCT on a log scale.Source: B. Performance under Failure-Free… · p. 6, Figure 3 · p. 7
  • On small and medium workloads at λ = 0, OAK’s tail JCT is 0–10% higher than Sia’s because the age key can reorder allocations to satisfy fairness when no job is at risk.Source: B. Performance under Failure-Free… · pp. 6–7

C. Failure Resilience and Tail Behaviour

Under failures, OAK improves JCT over Sia and over an aggregate restart estimator. Its ablation separates the restart factor’s mean-JCT gains from the age key’s protection against extreme jobs.

C. Failure Resilience and Tail Behaviour

  • At λ = 0.10, OAK reduces mean JCT over Sia by 15.2–70.9% and worst-case JCT by 33.5–75.8% across four workloads.Table II reports mean and tail JCT in seconds across four workloads and three failure rates.Source: Table II · p. 7, C. Failure Resilience and Tail Behaviour · p. 7
  • At λ = 0.10 on medium, OAK’s mean JCT is 21,740 versus Sia’s 75,730 and Pollux’s 17,814; Pollux’s mean JCT grows less across the sweep but uses continuous NSGA-driven reallocation.Across the full failure-rate sweep, Pollux’s mean JCT grows by 16% on medium and 15% on congestion, versus Sia’s 480% and 300%, respectively.Source: C. Failure Resilience and Tail Behaviour · p. 7
  • At λ = 0.10, OAK reduces mean JCT by 58.3% on medium and 33.9% on congestion; at λ = 0.15, it leads by 39% and 18%, respectively.Across the sweep, OAK and Sia are within 3% at λ ≤ 0.05; the gap narrows at λ = 0.15 as both struggle with extreme failure density.Source: C. Failure Resilience and Tail Behaviour · p. 7, Figure 4 · p. 8

D. Decomposed Restart Factor: Empirical Effect

  • At λ = 0.10, OAK’s decomposed factor cuts mean JCT 40% on ckpt-variance and 57% on medium versus the aggregate formula; cross-seed JCT standard deviation is 9–13× smaller.At λ = 0.05, the formulas produce mean JCT within 3% of each other.Source: Figure 5 · p. 8, D. Decomposed Restart Factor:… · p. 8
  • The aggregate estimate’s error grows with restarts: it over-penalizes workloads whose Prestart exceeds the measured 60 s per failure and under-penalizes those below 60 s.OAK uses measured Tckpt directly, so its decomposed factor remains accurate regardless of the parameterization.Source: D. Decomposed Restart Factor:… · p. 8
  • At λ = 0.10 on medium, restart-only reduces mean JCT by 58.4% over Sia versus 3.4% for age-only; adding age lowers the maximum JCT by 21.3%.The maximum drops from 222,096 s to 174,782 s, while cross-seed standard deviation of mean JCT falls from 1,423 s to 426 s.Source: D. Decomposed Restart Factor:… · p. 9
  • The terms play complementary roles: the restart factor improves mean and typical-tail JCT, while age protects against extreme jobs and seeds.On congestion at λ = 0.10, age lowers cross-seed standard deviation from 720 to 564, a 22% reduction.Source: D. Decomposed Restart Factor:… · p. 9

E. Hybrid OAK: NSGA II as the OAK Scheduler

Hybrid OAK places OAK’s composite score inside Pollux’s NSGA-II loop, replacing its aggregate restart factor. It offers a slower alternative to OAK’s ILP that better handles the high-failure regime described here.

a) Setup

  • Hybrid OAK uses OAK’s composite score inside Pollux’s NSGA-II loop in place of Pollux’s aggregate restart factor.Its setup used population = 20 and generations = 20, or 400 candidate evaluations per round.Source: E. Hybrid OAK: NSGA II as the OAK… · p. 9

b) Decision time vs JCT trade off

  • Hybrid OAK took ∼285 ms per round, compared with 7.65 ms for OAK and 9,909 ms for Pollux at its published defaults.The paper reports Hybrid OAK’s decision time as two orders of magnitude below Pollux’s default.Source: E. Hybrid OAK: NSGA II as the OAK… · p. 9
  • Hybrid OAK occupies 0.5% of a 60-second round while tracking Pollux-level JCT on medium and congestion across the failure-rate sweep.Source: E. Hybrid OAK: NSGA II as the OAK… · p. 9

c) Scheduler choice

F. Real-Hardware Validation

On a 4-V100 node, the real-hardware experiment reproduced OAK’s failure-resilience advantage over Sia and the schedulers’ decision-time ordering. It also found lower mean JCT for the NSGA-II schedulers than for the ILP schedulers in this small workload.

a) Setup

  • The experiment used four ResNet-18/CIFAR-10 jobs on one 4× V100 node, three seeds per cell, and Poisson failure rates from 0.0 to 0.10.The stated objective was to confirm directional ordering and decision-time differences on physical hardware, not reproduce absolute JCT.Source: F. Real-Hardware Validation · p. 10

b) Results

  • At λ = 0.10, OAK averaged 2.0 restarts versus Sia’s 4.7; the paper attributes this allocation pattern to the decomposed restart factor.Source: F. Real-Hardware Validation · p. 10
  • At λ = 0.10, OAK reduced mean JCT by 39.0% versus Sia (297.3 vs 487.6 s); tail JCT improved by 28.1% (511.8 vs 712.0 s).At λ = 0, OAK and Sia were near parity: 263.2 vs 257.8 s, a +2.1% difference.Source: Table IV · p. 10, F. Real-Hardware Validation · p. 10
  • Across the failure-rate sweep, Pollux’s mean JCT was 175.0/190.0/180.0 s versus Hybrid OAK’s 180.0/195.2/184.2 s, while Hybrid OAK had the lowest cross-seed standard deviation.The reported Hybrid OAK standard deviations ranged from 0.4 to 9.4.Source: F. Real-Hardware Validation · p. 10
  • Decision-time ordering persisted on hardware: ILP schedulers took ∼0.05 ms, Hybrid OAK ∼220 ms, and Pollux ∼1,420 ms per round.Hybrid OAK was within ∼3% of Pollux’s mean JCT at ∼6.4× tighter decision latency.Source: F. Real-Hardware Validation · p. 10

VI. THREATS TO VALIDITY

The paper reports improvements from OAK’s restart factor and a solver-agnostic hybrid, while qualifying the evidence by cluster scale, workload and failure model, and simulator abstractions. It also notes that the age key is a soft starvation guard, not a formal guarantee.

VI. THREATS TO VALIDITY

  • Scheduler comparisons used paired workloads, Poisson failure streams, trace-derived goodput, and seeds; reported JCT differences are invariant to absolute throughput calibration.Sensitivity to p, α, and Kmax was bounded below 1.5%.Source: §VI · p. 10
  • Evaluation covers 12-GPU simulation and 4-GPU hardware, data-parallel jobs, and Poisson failures; larger clusters and other workload or failure regimes remain untested.Source: §VI · p. 10
  • Most results are simulator-based and abstract over NCCL overhead, CPU-bound data pipelines, and node-level failure correlation; the 4-V100 validation is directional, not quantitative.Source: §VI · p. 10
  • The age key bounds priority inversion but does not guarantee starvation freedom; proofs under adversarial failures and fairness bounds under correlated node faults are left for future work.Source: §VI · pp. 10–11

VII. CONCLUSION

  • Hybrid OAK shows the composite utility can be used with reduced-population NSGA-II to retain reallocation robustness at high failure rates with lower per-round decision cost than Pollux’s default search.Source: §VII · p. 11
  • Future work includes scaling hardware validation beyond 4 GPUs, extending the utility to pipeline-parallel workloads, and studying Hybrid OAK’s population and iteration trade-off.Source: §VII · p. 11
Loading 2609.19024v2…