Papers
Topics
Authors
Recent
Search
2000 character limit reached

Dynamic Load Balancing for Uncertainty Quantification with Applications in Bayesian Inversion

Published 24 Jun 2026 in cs.DC | (2606.25693v1)

Abstract: Uncertainty Quantification (UQ) workflows present a particular scheduling challenge in high performance computing environments, as they typically generate large numbers of heterogeneous model evaluations with loose but non-trivial dependencies between tasks. A static one-size-fits-all approach in traditional schedulers is inadequate to handle heterogeneous tasks optimally. We introduce an improved load balancer in the UQ and Modelling Bridge (UM-Bridge) framework aimed at mitigating these issues; UM-Bridge is a language-agnostic interface developed to couple UQ software with numerical simulation. As a realistic example, we test the load balancer with a Bayesian inverse problem solved via multilevel delayed acceptance sampling. The underlying forward problem is a hierarchy of tsunami simulations enabled through ExaHyPE, whose runtimes span several orders of magnitude and loose dependencies between levels make the workload particularly challenging to schedule. Our results indicate the load balancer is effective at distributing the sampling requests with an average node idle time of close to a millisecond, while not making any prior assumptions about the workload.

Authors (3)

Summary

  • The paper introduces a persistent, language-agnostic UM-Bridge load balancer that uses SLURM job arrays and server pooling to eliminate per-request initialization without requiring runtime or dependency estimates.
  • The paper demonstrates average inter-request idle times of approximately 10⁻³ seconds while scheduling GP, coarse PDE, and fine PDE evaluations spanning six orders of magnitude in runtime.
  • The paper validates the approach on a Bayesian inversion of the 2011 Tōhoku earthquake, recovering posterior estimates consistent with an independent reference while identifying adaptive resource allocation and checkpointing as future needs.

This paper presents an improved load balancer for the UM-Bridge framework, designed to schedule heterogeneous uncertainty quantification (UQ) workloads on SLURM-managed HPC clusters. The authors demonstrate the approach on a Bayesian inverse problem for the 2011 Tōhoku earthquake source location, solved with multilevel delayed acceptance MCMC (MLDA) over a hierarchy of ExaHyPE tsunami simulations whose runtimes span six orders of magnitude (2606.25693).

Motivation and scheduling challenge

UQ workflows generate large numbers of model evaluations with heterogeneous runtimes and loose but non-trivial dependencies between tasks. The authors cite production data from Lawrence Livermore National Laboratory showing that 48.1% of jobs involve at least 100 identical jobs submitted in a short timeframe, indicating a shift away from the long-running homogeneous jobs that schedulers such as SLURM, PBS, and LoadLeveler were designed to handle optimally.

Existing remedies have drawbacks that motivate this work: HyperQueue coordinates task graphs atop native schedulers but is a separate tool; VECMA-ecosystem solutions interface with QCG-PilotJob or Dask but either require intrusive changes to application code, assume a priori knowledge of task runtimes, or are tied to a specific language or UQ ecosystem. The paper's contribution is a load balancer inside UM-Bridge—a language-agnostic HTTP client–server interface coupling UQ algorithms with simulation codes—that requires no admin privileges, no system-level changes, and no prior assumptions about task durations or dependencies.

The key deficiency addressed is per-request server initialisation in the two existing UM-Bridge load balancers [11018268]. While acceptable for long-running simulations, this cost dominates short tasks—specifically the millisecond-scale Gaussian process (GP) surrogate evaluations in the MLDA hierarchy used here—rendering those balancers unsuitable for heterogeneous workloads.

Load balancer design

The new implementation uses SLURM job arrays to request one bulk allocation at start-up and persistently manages a pool of UM-Bridge model servers for the entire run. Client requests are dispatched first-come, first-served in round-robin fashion across the pool. Concurrency is handled with a mutex and std::queue preserving arrival order; a std::conditional_variable lets threads sleep without polling, with notify_all() waking queued threads as soon as a server frees up. Because all data dependencies are managed by the UQ client, the balancer itself needs no knowledge of the algorithm, model structure, or expected runtimes—an implication being that it applies beyond UQ to any client–server workload where components are separable.

The forward model is abstracted as a map F:RnRm\mathcal{F}:\mathbb{R}^n \rightarrow \mathbb{R}^m, with optional exchange of Jacobians, gradients, and Hessians enabling gradient-based samplers such as HMC and NUTS through the same non-intrusive interface.

Forward solver: ExaHyPE and the Tōhoku configuration

ExaHyPE solves first-order hyperbolic balance laws using ADER-DG with an a posteriori subcell finite-volume limiter: a DG candidate step is accepted only if it satisfies physical admissibility and a relaxed discrete maximum principle; rejected cells are recomputed on $2p+1$ FV subcells per coordinate direction and projected back via a constrained conservative least-squares reconstruction. For the Tōhoku scenario, the limiter is deliberately geometry dominated: coastline, wet/dry, source-region, and boundary cells are always advanced by the FV layer, and a moving band of limiter cells follows the wavefront. Bathymetry is carried as a time-independent state component so that limiter projections act on the full state (h,hu,hv,b)(h, hu, hv, b); treating hh and bb independently across layers would destroy the lake-at-rest balance and trigger spurious limiting. An augmented Riemann solver balances the bed-slope term within the interface solve, enforcing non-negative depth and handling wet/dry interfaces.

Bayesian inverse problem and MLDA

The inverse problem recovers the location of the initial seafloor displacement from real bathymetry (GEBCO) and DART buoy data (probes 21418 and 21419), using a uniform prior on a [200,200]×[200,200][-200,200]\times[-200,200] km window and a Gaussian likelihood over wave height and arrival time at both probes. Inference uses three-level MLDA: level 2 is a fine-grid ExaHyPE solve, level 1 a coarse-grid ExaHyPE solve, and level 0 a GP surrogate trained on 512 Latin hypercube samples of level 1 output, implemented in PyTorch with a Matérn-$5/2$ kernel and automatic relevance determination.

Level Mean runtime [s] Cell size [m] DOF E[x]\mathbb{E}[x] E[y]\mathbb{E}[y]
0 (GP) 0.03 N/A 512 −204.45 −7962.52
1 (coarse PDE) 143.03 22197.50 656100 −6908.64 −2767.88
2 (fine PDE) 3071.53 7399.18 5904900 −2638.03 −10607.09

(Expectations are given in metres for both displacement coordinates.) A total of 1,500,005 level-0, 3,005 level-1, and 155 level-2 samples were evaluated. Variance decreases monotonically across levels, consistent with the telescoping estimator structure of MLDA, though the authors note the drop between levels 0 and 1 is larger than expected because the GP is conditioned on level-1 samples. The recovered posterior mean is consistent with the reference solution at the origin estimated independently from dynamic rupture modelling [10.1093/gji/ggu203], which validates the statistical component of the workflow alongside the computational one.

Load balancer performance

Experiments ran on the DINE2 partition of COSMA at Durham University (HDR200 InfiniBand, dual Intel Xeon Gold 6430), using five parallel MLDA chains in a 5-element job array, each task requesting 2 MPI ranks and 16 OpenMP threads, on non-exclusive nodes without special priority. Server uptime traces show densely packed bars with few gaps despite the six-order-of-magnitude runtime spread and strict sequential dependencies between levels. Measured idle times averaged O(103)\mathcal{O}(10^{-3}) s—consistent with HTTP communication overhead and substantially below the per-request initialisation cost of prior UM-Bridge balancers—with outliers up to roughly 0.1 s attributed to network congestion or dependency blocking. This result implies that even millisecond-scale GP evaluations can be co-scheduled with hour-long fine-grid PDE solves on shared infrastructure without dedicated tuning.

Limitations and open questions

Several constraints qualify these results. The evaluation covers a single cluster, a single application, and five chains; no comparison against HyperQueue, QCG-PilotJob, or Dask is reported quantitatively. The round-robin dispatch is workload-agnostic by design, so it cannot exploit knowledge of task runtimes or node utilisation—the authors identify node-utilisation-aware allocation as future work. There is no checkpointing mechanism, leaving resilience for lengthy workflows unaddressed. Whether the balancer sustains its idle-time performance under gradient-based MCMC methods, which impose additional heterogeneous demands (derivative evaluations), remains an open question the authors explicitly defer. Finally, the anomalous variance drop between levels 0 and 1 is observed but not analysed.

Conclusion

The paper delivers a deployable, assumption-free load balancer for UM-Bridge that eliminates per-request server initialisation via persistent SLURM job-array allocations. On a realistic Tōhoku tsunami inversion with MLDA, it achieves average inter-request idle times near one millisecond across a workload spanning six orders of magnitude in task duration, while recovering a posterior consistent with an independent reference solution. The remaining questions concern adaptive resource awareness, fault tolerance, and generalisation to derivative-heavy sampling algorithms.

Paper to Video (Beta)

No one has generated a video about this paper yet.

Whiteboard

No one has generated a whiteboard explanation for this paper yet.

Open Problems

We haven't generated a list of open problems mentioned in this paper yet.

Tweets

Sign up for free to view the 1 tweet with 0 likes about this paper.