arXiv CSopen access

Dynamic Load Balancing for Uncertainty Quantification with Applications in Bayesian Inversion

Dynamic Load Balancing for Uncertainty Quantification with Applications in Bayesian Inversion Chung Ming Loi1 , Mario Wille2 , and Anne Reinarz1

arXiv:2606.25693v1 [cs.DC] 24 Jun 2026

1

Durham University, UK {chung.m.loi,anne.k.reinarz}@durham.ac.uk 2 Technical University of Munich, Germany [email protected]

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. Keywords: High Performance Computing · Task Scheduling · Hyperbolic Partial Differential Equation · Multilevel Delayed Acceptance · Gaussian Process

1

Introduction

High Performance Computing (HPC) has enabled researchers to tackle once computationally intractable problems, including large-scale Bayesian inverse problems in which the posterior distribution over unknown parameters must be recovered from indirect observations. These problems arise naturally in geophysical applications such as seismology and tsunami modelling, where the forward model is an expensive numerical simulation of a Partial Differential Equation (PDE) [13,18,36]. Alongside advancements in the hardware landscape, scheduling and resource management constitute yet another factor to utilise HPC efficiently. Schedulers like SLURM, PBS, and LoadLeveler were designed to work optimally

2

C. M. Loi et al.

with traditional HPC workloads, where jobs are long-running and homogeneous. However, production clusters at Lawrence Livermore National Laboratory [1, 20] reported that 48.1% of jobs involve at least 100 identical jobs submitted in a short timeframe, signalling the shift in HPC workloads. Uncertainty Quantification (UQ) workflows are a particularly demanding example of this shift, as they require large numbers of model evaluations with heterogeneous runtimes and, in many cases, non-trivial dependencies between tasks that complicate scheduling further. Classical inverse problems deterministically apply the inverse map to find the best fit solution that describes the observation. However, they are notoriously ill-posed in the Hadamard sense: a solution may not exist, may not be unique, or may not depend continuously on the data [35]. Tikhonov regularisation [38] and related variational approaches can restore well-posedness, but yield only a point estimate with no quantification of uncertainty. The Bayesian formulation [22, 35] addresses both issues by treating the unknown parameters as a random variable, and conditioning it on the prior knowledge and observation data. It regularises the problem whilst recovering a full posterior distribution. A common approach to Bayesian inversion is Markov Chain Monte Carlo (MCMC), which constructs a Markov chain whose stationary distribution is the target posterior [35]. A large class of MCMC algorithms rely on sampling-based proposals, e.g. random walk Metropolis–Hastings [25], NUTS [21], and HMC [6], each requiring tens of thousands of forward model evaluations to converge, and some additionally require derivative information such as gradients or Hessians of the forward model. However, regardless of the proposal, MCMC faces a severe computational bottleneck in most applications: samples are proposed serially, each requiring one forward solve of the underlying numerical model, and Monte Carlo convergence is slow [14, 15]. Multilevel methods have emerged to address this bottleneck by exploiting a hierarchy of model approximations with increasing accuracy and cost [5]. In particular, we employ the Multilevel Delayed Acceptance (MLDA) algorithm [24] that generalises the Delayed Acceptance (DA) algorithm [4] by filtering proposed samples through a sequence of coarse models before committing to an expensive fine-level evaluation. While MLDA substantially improves sampling efficiency, the resulting workload is heterogeneous across levels and carries loose but non-trivial dependencies between proposals at adjacent levels, making it challenging to schedule efficiently on HPC systems. A growing number of tools have emerged to remedy the challenges in HPC scheduling. HyperQueue [2] is a plugin scheduler that works on top of the native scheduler. It is specifically designed to simplify execution of large workflows on HPC systems by coordinating a network of allocated nodes that execute submitted tasks. Self-contained UQ software ecosystems like the VECMA project [16] address the scheduling issues by interfacing with tools like QCG-PilotJob [3] and Dask [29], both of which are available as Python libraries that provide an API to interact with the native scheduler. However, these tools either require intrusive changes to the application code, assume a priori knowledge of task runtimes, or are tightly coupled to a specific programming language or UQ ecosystem, limiting their general applicability.

Dynamic Load Balancing for Uncertainty Quantification

3

In this work, we focus on workloads enabled by the UM-Bridge framework [31, 32]. It is a programming language agnostic interface designed to link UQ techniques and numerical simulations through the Hypertext Transfer Protocol (HTTP) client and server architecture, abstracting the forward model as a map F : Rn → Rm . Where available, the interface also supports the exchange of derivative information, including Jacobians, gradients, and Hessians of F, thus enabling gradient-based UQ algorithms such as HMC and NUTS to be coupled with simulation codes in the same non-intrusive way. By separating concerns between the UQ and simulation components, UM-Bridge allows non-experts to scale up their applications without modifying either the UQ or simulation code, and makes it straightforward to swap out either component independently. We introduce an improved load balancer in UM-Bridge for scaling up UQ workloads in SLURM-enabled HPC clusters that does not require system-level changes and makes no prior assumptions about task runtimes or dependencies. As a realistic example, we solve a Bayesian inverse problem following the example in [33]. Bayesian inversion for tsunami source parameters is an active research area, with recent work demonstrating the feasibility of large-scale inference [19, 30, 34]. We use the ExaHyPE engine to simulate the 2011 Tōhoku tsunami, which occurred due to an earthquake in the Japan trench. The forward model is a shallow water tsunami simulation of the depth-averaged shallow water equations implemented in ExaHyPE [28, 37]. Our inverse problem is then the task to recover the location of the earthquake from the tsunami produced by ExaHyPE and real ocean buoy data. Inference is performed using MLDA with a three-level model hierarchy: the two finest levels are ExaHyPE shallow water simulations of differing spatial resolutions, following the approach of [33], while the coarsest level replaces a further coarse PDE solve with a Gaussian Process (GP) surrogate [27] trained on fine-level simulation output, reducing the cost of the most frequently evaluated level by several orders of magnitude. The three-level model hierarchy, whose runtimes span several orders of magnitude, combined with the loose dependencies between levels inherent to MLDA, provides a demanding test of the load balancer. The paper is structured as follows. The load balancer design and implementation are presented in Section 2. Section 3 introduces ExaHyPE as a framework for simulating first-order hyperbolic PDEs and describes the Tōhoku tsunami configuration. Section 4 formulates the Bayesian inverse problem, and Section 5 describes the MLDA algorithm. Finally, Section 6 presents numerical results demonstrating both the efficacy of the load balancer and the recovered posterior distribution for the tsunami source location, before we conclude in Section 7.

2

SLURM-Integrated Load Balancing in UM-Bridge

We introduce an improved UM-Bridge load balancer that directly addresses the scheduling challenges of MLDA workloads identified in Section 1, in particular the elimination of per-request initialisation overhead that made existing implementations [23] unsuitable for heterogeneous task durations.

4

2.1

C. M. Loi et al.

Architecture Overview

UM-Bridge [31, 32] is a language-agnostic interface that connects UQ algorithms with numerical models non-intrusively via an HTTP client-server architecture. The forward model is abstracted as a map F : Rn → Rm , evaluated at points {θ i }N i=0 that may be predetermined or determined adaptively, as in Bayesian inversion. Where available, UM-Bridge also supports the exchange of derivative information—Jacobians, gradients, and Hessians of F—enabling gradient-based MCMC algorithms such as HMC and NUTS to be coupled with simulation codes in the same non-intrusive way. By treating UQ algorithms and numerical solvers as entirely separate applications, UM-Bridge allows components written in different programming languages to be coupled and swapped out without modification to either side, lowering the barrier to non-experts scaling up their applications. The load balancer routes evaluation requests {F(θ i )}N i=0 across available servers; since all data dependencies are managed by the UQ client, the load balancer requires no knowledge of the UQ algorithm, model structure, or expected runtimes, making it broadly applicable beyond UQ to any setting where client and server components are separable.

Shared Filesystem

Parallel UQ Client

UM

UM

Load Balancer

SLURM

UM Model MPI, Launcher

LAPTOP/LOGIN NODE/COMPUTE NODE

SINGULARITY/SLURMD

MPI 1..*

Model MPI, Worker SINGULARITY/SLURMD COMPUTE CLUSTER

Fig. 1: Load balancer configuration for multiple instances of UM-Bridge models and a parallel client.

2.2

SLURM-integrated Load Balancer

UQ workflows based on sampling create a distinctive scheduling problem: they submit large numbers of heterogeneous tasks in a short timeframe [10], with runtimes that are often unpredictable and span several orders of magnitude. Traditional SLURM-based approaches handle this poorly for two reasons. First, each job submission incurs launch overhead that can dominate the runtime of short-lived tasks. Second, users must specify a time limit per job, which is typically set conservatively to the worst-case runtime, leading to systematic underutilisation when most tasks are completed much earlier. For MLDA workloads, these problems are compounded by the strict sequential dependencies between levels: a fine-level evaluation cannot proceed until the coarse filter has accepted a proposal, so idle time at one level directly delays progress at the next. The UM-Bridge load balancer is an intermediate abstraction layer between the parallel UQ client and model servers running on HPC compute nodes. Functionally, it forwards client requests to available servers, operating like a network proxy.

Dynamic Load Balancing for Uncertainty Quantification

5

Algorithm 1 Simplified algorithm describing the UM-Bridge load balancer. 1: parallel for j = 0 to N − 1 do ▷ Concurrent incoming requests 2: mutex.lock() 3: queue.push(request[j]) 4: if free server exists then ▷ Point of entry after wakeup 5: server ← getFreeServer() 6: request[j] ← queue.pop() 7: server.markBusy() 8: mutex.unlock() 9: return server(request[j]) ▷ Blocking call; reset busyness once done 10: else 11: conditional_variable.wait(mutex) ▷ Sleep current thread 12: go to 4 ▷ Triggered whenever a server is marked as free 13: end if 14: end parallel for

The two existing UM-Bridge load balancers [23] address job submission overhead partially, but both initialise a new UM-Bridge server per client request. This perrequest initialisation cost is acceptable for long-running simulations, but becomes the dominant overhead for short tasks such as the GP surrogate evaluations in our MLDA hierarchy, rendering them unsuitable for workloads with heterogeneous task durations. Our new implementation eliminates this bottleneck by using the SLURM job array functionality to request one bulk allocation at initialisation and persistently manage a pool of UM-Bridge model servers for the entire duration of the run, as shown in Fig. 1. Client requests are distributed across the pool in a first-come, first-served, round-robin fashion, with no assumption made about task duration or inter-task dependencies. The persistent server pool means that even millisecondduration GP evaluations are dispatched and returned with overhead limited to HTTP communication latency, while the same infrastructure simultaneously handles hour-long fine-grid PDE solves. Concurrent requests are handled safely using a mutex lock and std::queue to preserve arrival order and prevent race conditions. A std::conditional_variable allows threads to sleep and wake without polling: whenever a server completes a request and is marked free, a notify_all() call wakes all queued threads and the next request is dispatched immediately. This design keeps the average inter-request idle time at the scale of HTTP communication overhead, as confirmed by the experimental results in Section 6. Algorithm 1 describes the dispatch logic. The implementation requires no admin privileges and does not modify the native scheduler, making it deployable on any SLURM-enabled cluster without system-level changes. It is also compatible with existing UM-Bridge setups without modification.

6

C. M. Loi et al.

3

ExaHyPE: Exascale Hyperbolic PDE Engine

In this work we use ExaHyPE to simulate the 2011 Tōhoku tsunami. We briefly describe the numerical scheme and provide details on the tsunami application in the following sections. 3.1

ADER-DG and A Posteriori Finite-Volume Subcell Limiting

ExaHyPE3 is a simulation framework for first-order hyperbolic balance laws with conservative fluxes, non-conservative products, and source terms, ∂t Q +

D X d=1

∂xd Fd (Q) +

D X

Bd (Q) ∂xd Q = S(Q, x, t),

Q ∈ Rm .

(1)

d=1

using the engine design described in [28]. The regular discretization is an ADERDG method. On every cell K, the state is a tensor-product polynomial Qh |K ∈ Vhp (K). A local space-time predictor evolves this polynomial inside the cell, and a corrector couples neighbouring cells through face fluxes and non-conservative interface terms. This yields a high-order one-step method with low numerical diffusion in smooth solution regions [7, 8]. High-order polynomials are not monotone near shocks, discontinuities, steep gradients, or under-resolved features. Hence, the method uses a posteriori limiting: an ADER-DG step first proposes Q∗,n+1 , and this candidate is accepted only if h it satisfies physical admissibility and a relaxed discrete maximum principle [9]. For an observable w(Q), the maximum principle compares the candidate with extrema from the previous accepted state on K and its face-neighbourhood N (K), n mnN − δ ≤ w(Q∗,n+1 (x)) ≤ MN + δ, h

n δ = max(ϵ, α(MN − mnN )) .

(2)

The admissibility test enforces problem-specific invariant-domain constraints, such as finite states and positivity of selected state components. Accepted cells remain on the ADER-DG layer; rejected cells are recomputed on a robust subcell FV layer. For a degree-p ADER-DG polynomial, the limiter uses Ns = 2p+1 subcells per coordinate direction. This samples all polynomial modes and keeps the FV and ADER-DG time-step restrictions compatible [9]. If ωi ⊂ K is an FV subcell, the rollback state is projected from the accepted DG polynomial to subcell averages, Z 1 n Q̄i = Qn (x) dx, Q̄n = ΠDG→FV Qnh . (3) |ωi | ωi h The FV method recomputes the same time interval on these subcells. The accepted subcell solution is then projected back to Vhp (K) by a constrained conservative least-squares reconstruction [9]. Writing N = p for the DG polynomial degree, 3

Available as Docker image at https://hub.docker.com/r/peanoframework/exahype2

Dynamic Load Balancing for Uncertainty Quantification

7

Ns = 2N + 1 FV subcell values are used per coordinate direction to recover only N + 1 DG coefficients. The FV-to-DG reconstruction is therefore overdetermined, so a constrained least-squares reconstruction is used, with conservation enforced through a Lagrange multiplier: Z X n+1 Qn+1 = Π Q̄ , Qn+1 dx = |ωi |Q̄n+1 . (4) FV→DG i h h K

i

The resulting method is locally hybrid: high-order ADER-DG is retained where the solution is smooth, while the FV subcell update supplies robustness near shocks, discontinuities, and other troubled regions. Crucially, the limiter acts only where the a posteriori admissibility check fails: the computational cost of the FV recomputation is therefore confined to a small fraction of cells at any given time step, preserving the efficiency of the high-order scheme globally. Figure 2 summarizes the rollback, subcell update, and conservative recovery between the two representations.

ADER-DG cell

polynomial state Qh ∈ Vhp (K)

troubled candidate is rolled back ΠDG→FV

ΠFV→DG robust FV update then conservative recovery

FV subcells

cell averages Q̄i , Ns = 2p + 1

Fig. 2: Schematic of the a posteriori subcell limiter. A rejected ADER-DG candidate is replaced by an FV recomputation on subcell averages and then projected back to the DG polynomial space.

3.2

2011 Tōhoku Tsunami

The tsunami is modelled with the depth-averaged shallow-water equations over a variable bathymetry b(x). We use the basic shallow-water equations with bathymetry source terms and neglect bottom friction as well as non-hydrostatic corrections [33]. With water depth h, horizontal velocity (u, v), momenta (hu, hv), gravity g, and free surface η = h + b, the system is written as       h hu hv 0 hu  hu2 huv  gh ∂x (b + h)      ∂t  (5) hv  + ∇· huv hv 2  + gh ∂y (b + h) = 0. b 0 0 0 The bathymetry is carried as a time-independent state component, with the pressure gradient and bed slope represented together through gh∇(b + h),

8

C. M. Loi et al.

exposing the lake-at-rest balance (u, v) = 0, η = const. that the numerical method must preserve while allowing large bathymetry jumps, dry land, and inundation. The Tōhoku configuration uses the domain [−499, 1299] × [−949, 849] km around Japan, with the ocean initially at rest and the earthquake displacement introduced as a filtered change of bed elevation. The plotted quantity is the sea-surface-height anomaly (SSHA), i.e. the deviation of η from the initial still-water state. Although b has zero physical time evolution, it is carried as an unknown in both the ADER-DG and FV layers so that limiter projections act on the full state (h, hu, hv, b). This is required for consistency: since η = h + b and the bed-slope source depend on h and b together, treating them independently across layers would destroy the lake-at-rest balance, introduce artificial SSHA, and trigger spurious limiting. The FV subcell layer uses an augmented-state Riemann solver [12] that decomposes a jump augmented with the momentum flux and bathymetry step, balancing the bed-slope term inside the interface solve rather than as a separate correction. This preserves well-balancedness over rough bathymetry, enforces non-negative water depth, and handles wet/dry interfaces as inundation problems, with one-sided draining capped so that no FV update removes more water than there is locally available.

800

3

t = 340.9 s

t = 1363.9 s

t = 3068.7 s

600

2

400

y [km]

0

0

−200

SSHA [m]

1

200

−1

−400 −600

−2

−800

−3 800 600 400

y [km]

200

Regular cells 0 −200

Limiter cells

−400 −600 −800 −400 −200

0

200

400

x [km]

600

800 1000 1200

−400 −200

0

200

400

x [km]

600

800 1000 1200

−400 −200

0

200

400

600

800 1000 1200

x [km]

Fig. 3: Tōhoku tsunami snapshots at t = 340.9 s, 1363.9 s, and 3068.7 s. Top: SSHA at these times. Bottom: cells treated by the finite-volume limiter for the same times.

For this scenario, the limiter is deliberately geometry dominated. A DG cell remains on the ADER-DG layer only if the state is finite, the water column exceeds the wet/dry threshold, the cell lies in deep water, and the cell-local variation of η is small relative to the local depth; coastline, wet/dry, source-region, and

Dynamic Load Balancing for Uncertainty Quantification

9

boundary cells are therefore advanced through the FV limiter. Boundary cells are not physically troubled but are limited because the a posteriori troubledness test is underdetermined there: no exterior DG neighbour polynomial exists, so a one-cell FV layer is used to remove one-sided reconstruction artefacts while applying the same positivity-preserving Riemann treatment used near land and shallow water. Figure 3 visualizes three snapshots. The upper row shows SSHA, while the lower row shows limiter ownership. The open-ocean wave is advanced by the regular ADER-DG scheme; limiter cells are concentrated where the posteriori checks are intentionally conservative. In addition to the stationary coastline and boundary layer, a moving band of limiter cells follows the tsunami wavefront, indicating where the evolving free-surface gradients temporarily require the robust FV subcell update.

4

The Bayesian Inverse Problem

A Bayesian inverse problem is the task of reconstructing the posterior distribution P(θ | y) of uncertain parameters θ given measurement data y and a forward model F. By Bayes’ theorem, P(θ | y) ∝ L(y | θ) π0 (θ),

(6)

where L(y | θ) is the likelihood of observing y given θ, and π0 (θ) is the prior distribution encoding knowledge of θ before observing data. The normalising constant P(y) is intractable in all but the simplest cases and cancels in MCMC acceptance ratios, so only the proportional form (6) is required in practice.

800

600

400

y [km]

200

0

−200

−400

−600

−800

−400

−200

0

200

400

600

800

1000

1200

x [km]

Fig. 4: Tōhoku computational domain used for the uncertainty-quantification workflow. The red box marks the displacement translation window [−200, 200] × [−200, 200] km.

In this work, we are concerned with the 2011 Tōhoku tsunami, which resulted from an earthquake in the Japan trench. Our inverse problem is recovering

10

C. M. Loi et al.

the location of the initial displacement from real-world bathymetry4 and ocean buoy5 data. We choose π0 (θ) as a 2D uniform distribution to reflect zero a priori knowledge about the unknown coordinates with support chosen to keep the initial displacement away from the domain boundary and sufficiently far from the probes to avoid unphysical simulation output. Figure 4 shows the prior on the computational domain. We formulate the inverse problem in accordance with [33], where the Gaussian likelihood is chosen to quantify the mismatch between observation and simulation. The mean vector contains the wave height and arrival time of the tsunami at the two DART probes 21418 and 21419, while the diagonal covariance matrix encodes both measurement noise in the probes and the discrepancy between the numerical model and reality.

5

Multilevel Delayed Acceptance MCMC

MCMC methods construct a Markov chain on the parameter space such that the chain converge to the target distribution iteratively, thus allowing posterior quantities of interest to be estimated from the resulting samples. The MetropolisHastings (MH) algorithm [17,26] is the canonical method to construct such chains, but its efficiency is heavily dependent on the choice of proposal distribution. Many alternatives have been proposed to increase sampling efficiency. In this section, we describe the classic DA algorithm proposed by [4] and its multilevel extension due to [24]. 5.1

Delayed Acceptance MCMC

The original DA algorithm is a two-level method that aims to increase the sampling efficiency. It introduces a step that filters out bad proposal early using a coarse-model before the expensive evaluation at the fine-level. As such, the fine-level has higher accuracy but at the cost of time to solution. Conversely, the coarse-level is relatively inaccurate but fast to compute. Common choices for the coarse-level include surrogate models developed using a collection of simulation outputs, or simply the fine-model with a coarser computational domain. Algorithm 2 describes the DA sampler. Rather than evaluating each proposal directly with the expensive fine-model, a preliminary MH step targeting the cheap coarse density πC (·) filters out poor proposals before they reach the fine-level. The surviving state ψ is accepted or rejected at the fine-level with probability αF (ψ | θj ), which corrects for the discrepancy between the coarse and fine densities, and preserves the correct stationary distribution. The key computational saving arises because any proposal rejected at the coarse stage never triggers a fine-model evaluation, yielding a net reduction in the number of expensive forward solves relative to standard MH that grows with the rejection rate of the coarse-level. 4 5

https://www.gebco.net/ https://www.ndbc.noaa.gov/

Dynamic Load Balancing for Uncertainty Quantification

11

Algorithm 2 The outline of the DA sampler

1: function {θ1 , . . . , θN } = DA πF , πC , q, θ0 , N 2: input: fine target πF (·); coarse target πC (·); proposal q(·|·); initial state θ0 ; number of steps N 3: output: chain θ1 , . . . , θN targeting πF 4: for j = 0 to N − 1 do
5: ψ = MH πC , q, θj , 1   πF (ψ) πC (θj ) j j+1 6: Set θ = ψ with probability αF (ψ | θ ) = min 1, , otherwise πF (θj ) πC (ψ) j+1 j θ =θ 7: end for

5.2

Multilevel Generalisation

While the DA algorithm offers better sampling efficiency over standard MH, the acceptance rate at the fine-level is fundamentally limited by how well the coarsemodel approximates the fine density. Multilevel methods address this by defining a hierarchy of models with increasing fidelity, so that each level proposes to the next using a progressively better approximation of the target. A proposed sample passes through several coarse filters before reaching the expensive fine-model, with each filter improving the quality of proposals seen at the next level. This telescoping structure yields variance reduction in the estimators at finer levels, as the coarse levels absorb the bulk of the exploration cost whilst passing only well-informed proposals upward.

Fig. 5: Two level accept-reject scheme of the MLDA algorithm.

Figure 5 pictorially describes a two level MLDA chain, which generalises DA by replacing the single coarse step with a randomised subchain of length

12

C. M. Loi et al.

nℓ ∼ pℓ (·), allowing different state spaces at adjacent levels, and generating each subchain recursively via MLDA at level ℓ − 1, reverting to MH at ℓ = 0. The expectation of the quantity of interest ϕ from MLDA is given by a telescoping sum EπL [ϕL ] = Eπ0 [ϕ0 ] +

L X

Eπℓ [ϕℓ ] − Eπℓ−1 [ϕℓ−1 ] ,

(7)

ℓ=1

where Eπℓ [·] denotes expectation with respect to the posterior at level ℓ.

6

Tsunami Inversion and Load Balancer Performance

All simulations ran on the DINE2 partition of the COSMA cluster at Durham University. Nodes are connected via HDR200 InfiniBand, each equipped with two Intel Xeon Gold 6430 32-core processors and 2 TB of RAM. We request 2 MPI ranks and 16 OpenMP threads per rank for each task in a 5-element job array, hosting 5 parallel MLDA chains. The node is not exclusive and no special priority was requested, reflecting a realistic production environment. MLDA sampling is facilitated by the tinyDA Python package6 , showcasing UM-Bridge’s ability to couple UQ and simulation software written in different programming languages. All code to reproduce these experiments is available at a GitHub7 repository. Level ℓ t̄ℓ [s] hℓ [m] DOF E[ϕℓ ] V[ϕℓ ] 0 0.03 N/A 512 -204.45 -7962.52 20329.02 43285.31 1 143.03 22197.50 656100 -6908.64 -2767.88 13279.11 28717.72 2 3071.53 7399.18 5904900 -2638.03 -10607.09 8680.50 15659.04 Table 1: MLDA model hierarchies and inverse problem results. The hℓ column lists the cell size of the computational domain. The Degrees of Freedom (DOF) at level 0 is defined as the dimension of the kernel matrix, whereas the remaining two states the DOF in the DG solver. The expectation (in metres) and variance are given for both coordinates of the initial displacement: x and y.

6.1

Inverse Problem: Tōhoku Tsunami

The MLDA algorithm is configured with three levels: level 2 is a fine grid ExaHyPE model, level 1 is a coarse grid ExaHyPE model, and level 0 is a GP surrogate trained on 512 Latin Hypercube samples drawn from the level 1 model. The GP is implemented in PyTorch with a 5/2 Matérn kernel, zero mean, and automatic relevance determination; hyperparameters are optimised by maximising 6 7

https://github.com/mikkelbue/tinyDA/tree/main https://github.com/chun9l/ExaHyPE2_UQ

Dynamic Load Balancing for Uncertainty Quantification

13

the marginal likelihood on the training data. The O(n3 ) cost of kernel matrix operations is negligible at n = 512 relative to the PDE solves at finer levels. The MLDA setup is summarised in Table 1. The runtimes vary by orders of magnitude, with each level averaging 0.03 s, 143.03 s, and 3071.53 s respectively. The loose dependencies between levels inherent to MLDA, where finer samples depend on the acceptance of coarse samples, make this workload particularly challenging to schedule in an HPC environment. In total, the algorithm evaluated 1,500,005 level 0, 3,005 level 1, and 155 level 2 samples. Figure 6 shows raw data from probe 21418 overlaid with level 0 draws from the prior and posterior distribution, using a separate GP trained to reconstruct the full time series.

Fig. 6: Raw time series data from probe 21418 and samples at level 0. Top: 50 sample draws from the uniform prior. Bottom: 50 sample draws from the recovered posterior.

Figure 7 illustrates the resulting sample densities in each MLDA levels. The dashed lines mark the posterior mean E[Qℓ ], and the red cross indicates the reference solution at the origin estimated by [11]. The primary goal of multilevel methods is to improve acceptance rates through a series of cheap filters, which are then corrected progressively with each acceptance at the finer level. Typically, this effect is reflected by the reduction in variance due to the telescoping sum (7). The expectation and variance are listed in Table 1. Indeed, we observe variance reduction across levels, notably the drop between levels 0 and 1 is larger than expected, since the GP is conditioned on level 1 samples. 6.2

UM-Bridge Load Balancer

The timings were obtained by recording the arrival and departure times of individual requests in the Python server script. Since data dependencies are managed by the tinyDA client, and the load balancer greedily assigns parallel requests to available UM-Bridge servers, the timestamps directly represent time spent inside a server and—by implication—the delay between requests. Figure 8 shows the uptime of each UM-Bridge server as a bar chart, with colours indicating model fidelity: green for level 2, orange for level 1, and blue for level 0. The three-level hierarchy presents a demanding scheduling problem: the level 0 GP evaluations are six orders of magnitude faster than the level 2 ExaHyPE runs, yet carry strict sequential dependencies within each MLDA chain; finer level evaluations cannot proceed until the coarse filter has accepted a proposal. The level 0 bars are not individually visible at this scale due to their negligible runtime of 0.03 s, but their dependencies are indicated by the arrows, which trace the request ordering within a single MLDA chain. The bars are

14

C. M. Loi et al.

Level 0

Level 1

Level 2

y [km]

100

0

−100

−50

0

x [km]

50

100

−50

0

x [km]

50

100

−50

0

50

100

x [km]

Fig. 7: Density of posterior samples from each MLDA levels. The dashed lines indicate the sample mean, and the red cross provides a known reference at (0, 0).

Fig. 8: The execution of client requests in the 5-element job array is indicated by the y-axis labels. The x-axis is the elapsed time. The colour of the bars indicates different model fidelities: green is level 2, orange is level 1, and blue is level 0. Arrows indicate the request dependencies within one MLDA chain; the numbers are reduced to avoid clutter. A higher-resolution image is available on the linked author’s GitHub repository.

densely packed with few whitespaces, indicating that despite task heterogeneity the loosely coupled MLDA requests were handled promptly across all five parallel chains.

Dynamic Load Balancing for Uncertainty Quantification

15

Fig. 9: Boxplot of idle time between sampling requests.

The idle times, shown in Fig. 9, averaged O(10−3 ) seconds, consistent with the overhead of HTTP communication in UM-Bridge and substantially lower than the per-request server initialisation overhead incurred by the existing UM-Bridge load balancers [23], which made them ineffective for short-running tasks such as the level 0 GP evaluations here. Outliers extending to approximately 0.1 s are attributed to occasional network congestion or requests blocked by pending the completion of a dependency. Taken together, these results demonstrate that the load balancer effectively handles a workload—with non-trivial inter-task dependencies—spanning six orders of magnitude in task duration without any prior knowledge of task runtimes, thus resolving issues with the existing load balancer implementation.

7

Conclusion

We have presented an improved SLURM-integrated load balancer for the UMBridge framework, targeting the heterogeneous and loosely dependent workloads that arise in UQ workflows on HPC systems. Tested on a Bayesian inverse problem for the 2011 Tōhoku earthquake source location using MLDA with a hierarchy spanning six orders of magnitude in runtime, the load balancer achieves average idle times of O(10−3 ) seconds across five parallel chains without prior knowledge of task runtimes or dependencies, substantially improving on existing UM-Bridge load balancers [23]. The recovered posterior is consistent with the known reference solution [11], validating both the statistical and computational components of the workflow. Future work includes incorporating node utilization awareness for more dynamic resource allocation, adding a checkpointing mechanism for resilience in lengthy workflows, and evaluating the load balancer on gradient-based MCMC methods that place additional heterogeneous demands on the scheduler. Acknowledgments. This work has made use of the Hamilton HPC Service and the COSMA supercomputer of Durham University. The COSMA supercomputer is

16

C. M. Loi et al.

a DiRAC@Durham facility managed by the Institute for Computational Cosmology on behalf of the STFC DiRAC HPC Facility (www.dirac.ac.uk). The equipment was funded by BEIS capital funding via STFC capital grants ST/K00042X/1, ST/P002293/1, ST/R002371/1 and ST/S002502/1, Durham University and STFC operations grant ST/R000832/1. DiRAC is part of the National e-Infrastructure. Disclosure of Interests. The authors have no competing interests to declare that are relevant to the content of this article.

References 1. Ahn, D.H., Bass, N., Chu, A., Garlick, J., Grondona, M., Herbein, S., Koning, J., Patki, T., Scogland, T.R.W., Springmeyer, B., Taufer, M.: Flux: Overcoming scheduling challenges for exascale workflows. In: 2018 IEEE/ACM Workflows in Support of Large-Scale Science (WORKS). pp. 10–19 (2018). https://doi.org/10. 1109/WORKS.2018.00007 2. Beránek, J., Böhm, A., Palermo, G., Martinovič, J., Jansík, B.: Hyperqueue: Efficient and ergonomic task graphs on hpc clusters. SoftwareX 27, 101814 (2024). https: //doi.org/https://doi.org/10.1016/j.softx.2024.101814, https://www.sciencedirect. com/science/article/pii/S2352711024001857 3. Bosak, B., Piontek, T., Karlshoefer, P., Raffin, E., Lakhlili, J., Kopta, P.: Verification, validation and uncertainty quantification of large-scale applications with qcg-pilotjob. In: Paszynski, M., Kranzlmüller, D., Krzhizhanovskaya, V.V., Dongarra, J.J., Sloot, P.M. (eds.) Computational Science – ICCS 2021. pp. 495–501. Springer International Publishing, Cham (2021) 4. Christen, J.A., Fox, C.: Markov chain monte carlo using an approximation. Journal of Computational and Graphical Statistics 14(4), 795–810 (2005). https://doi.org/ 10.1198/106186005X76983, https://doi.org/10.1198/106186005X76983 5. Dodwell, T.J., Ketelsen, C., Scheichl, R., Teckentrup, A.L.: Multilevel markov chain monte carlo. SIAM Review 61(3), 509–545 (2019). https://doi.org/10.1137/ 19M126966X, https://doi.org/10.1137/19M126966X 6. Duane, S., Kennedy, A., Pendleton, B.J., Roweth, D.: Hybrid Monte Carlo. Physics Letters B 195(2), 216–222 (Sep 1987). https://doi.org/10.1016/0370-2693(87) 91197-X, https://linkinghub.elsevier.com/retrieve/pii/037026938791197X 7. Dumbser, M., Balsara, D.S., Toro, E.F., Munz, C.D.: A unified framework for the construction of one-step finite volume and discontinuous Galerkin schemes on unstructured meshes. Journal of Computational Physics 227, 8209–8253 (2008) 8. Dumbser, M., Fambri, F., Tavelli, M., Bader, M., Weinzierl, T.: Efficient implementation of ADER discontinuous Galerkin schemes for a scalable hyperbolic PDE engine. Axioms 7(3), 63 (2018). https://doi.org/10.3390/axioms7030063 9. Dumbser, M., Loubère, R.: A simple robust and accurate a posteriori sub-cell finite volume limiter for the discontinuous Galerkin method on unstructured meshes. Journal of Computational Physics 319, 163–199 (2016) 10. Farcas, I.G., Merlo, G., Jenko, F.: A general framework for quantifying uncertainty at scale. Communications Engineering 1 (2022), https://api.semanticscholar.org/ CorpusID:246652492 11. Galvez, P., Ampuero, J.P., Dalguer, L.A., Somala, S.N., Nissen-Meyer, T.: Dynamic earthquake rupture modelled with an unstructured 3-d spectral element method applied to the 2011 m9 tohoku earthquake. Geophysical Journal International

Dynamic Load Balancing for Uncertainty Quantification

17

198(2), 1222–1240 (08 2014). https://doi.org/10.1093/gji/ggu203, https://doi.org/ 10.1093/gji/ggu203 12. George, D.L.: Augmented Riemann solvers for the shallow water equations over variable topography with steady states and inundation. Journal of Computational Physics 227(6), 3089–3113 (2008). https://doi.org/10.1016/j.jcp.2007.10.027 13. Gibbons, S.J., Lorito, S., Macías, J., Løvholt, F., Selva, J., Volpe, M., SánchezLinares, C., Babeyko, A., Brizuela, B., Cirella, A., Castro, M.J., de la Asunción, M., Lanucara, P., Glimsdal, S., Lorenzino, M.C., Nazaria, M., Pizzimenti, L., Romano, F., Scala, A., Tonini, R., Manuel González Vida, J., Vöge, M.: Probabilistic tsunami hazard analysis: High performance computing for massive scale inundation simulations. Frontiers in Earth Science Volume 8 - 2020 (2020). https://doi. org/10.3389/feart.2020.591549, https://www.frontiersin.org/journals/earth-science/ articles/10.3389/feart.2020.591549 14. Giles, M.B.: Multilevel monte carlo path simulation. Operations research 56(3), 607–617 (2008) 15. Giles, M.B.: Multilevel monte carlo methods. Acta numerica 24, 259–328 (2015) 16. Groen, D., Arabnejad, H., Jancauskas, V., Edeling, W.N., Jansson, F., Richardson, R.A., Lakhlili, J., Veen, L., Bosak, B., Kopta, P., Wright, D.W., Monnier, N., Karlshoefer, P., Suleimenova, D., Sinclair, R., Vassaux, M., Nikishova, A., Bieniek, M., Luk, O.O., Kulczewski, M., Raffin, E., Crommelin, D., Hoenen, O., Coster, D.P., Piontek, T., Coveney, P.V.: Vecmatk: a scalable verification, validation and uncertainty quantification toolkit for scientific simulations. Philosophical Transactions of the Royal Society A: Mathematical, Physical and Engineering Sciences 379(2197), 20200221 (2021). https://doi.org/10.1098/rsta.2020.0221, https://royalsocietypublishing.org/doi/abs/10.1098/rsta.2020.0221 17. Hastings, W.K.: Monte carlo sampling methods using markov chains and their applications. Biometrika 57(1), 97–109 (04 1970). https://doi.org/10.1093/biomet/ 57.1.97, https://doi.org/10.1093/biomet/57.1.97 18. Henneking, S., Venkat, S., Dobrev, V., Camier, J., Kolev, T., Fernando, M., Gabriel, A.A., Ghattas, O.: Real-time bayesian inference at extreme scale: A digital twin for tsunami early warning applied to the cascadia subduction zone. In: Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis. p. 60–71. SC ’25, Association for Computing Machinery, New York, NY, USA (2025). https://doi.org/10.1145/3712285.3771787, https://doi.org/ 10.1145/3712285.3771787 19. Henneking, S., Venkat, S., Ghattas, O.: Goal-oriented real-time bayesian inference for linear autonomous dynamical systems with application to digital twins for tsunami early warning. Journal of Computational Physics 552, 114682 (2026). https:// doi.org/https://doi.org/10.1016/j.jcp.2026.114682, https://www.sciencedirect.com/ science/article/pii/S002199912600032X 20. Herbein, S., Patki, T., Ahn, D.H., Mobo, S., Hathaway, C., Caíno-Lores, S., Corbett, J., Domyancic, D., Scogland, T.R., de Supinski, B.R., Taufer, M.: An analytical performance model of generalized hierarchical scheduling. The International Journal of High Performance Computing Applications 36(3), 289–306 (2022). https://doi. org/10.1177/10943420211051039, https://doi.org/10.1177/10943420211051039 21. Hoffman, M.D., Gelman, A., et al.: The No-U-Turn sampler: adaptively setting path lengths in Hamiltonian Monte Carlo. J. Mach. Learn. Res. 15(1), 1593–1623 (2014) 22. Latz, J.: Bayesian inverse problems are usually well-posed. SIAM Review 65(3), 831–865 (2023). https://doi.org/10.1137/23M1556435, https://doi.org/10.1137/ 23M1556435

18

C. M. Loi et al.

23. Loi, C.M., Reinarz, A., Seelinger, L., Hornsby, W., Buchanan, J., Lykkegaard, M.: A performance analysis of task scheduling for uq workflows on hpc systems. In: ISC High Performance 2025 Research Paper Proceedings (40th International Conference). pp. 1–14 (2025). https://doi.org/10.23919/ISC.2025.11018268 24. Lykkegaard, M.B., Dodwell, T.J., Fox, C., Mingas, G., Scheichl, R.: Multilevel delayed acceptance mcmc. SIAM/ASA Journal on Uncertainty Quantification 11(1), 1–30 (2023). https://doi.org/10.1137/22M1476770, https://doi.org/10.1137/ 22M1476770 25. Metropolis, N., Rosenbluth, A.W., Rosenbluth, M.N., Teller, A.H., Teller, E.: Equation of state calculations by fast computing machines. The Journal of Chemical Physics 21(6), 1087–1092 (06 1953). https://doi.org/10.1063/1.1699114, https: //doi.org/10.1063/1.1699114 26. Metropolis, N., Rosenbluth, A.W., Rosenbluth, M.N., Teller, A.H., Teller, E.: Equation of state calculations by fast computing machines. The Journal of Chemical Physics 21(6), 1087–1092 (06 1953). https://doi.org/10.1063/1.1699114, https: //doi.org/10.1063/1.1699114 27. Rasmussen, C., Williams, C.: Gaussian Processes for Machine Learning. Adaptative computation and machine learning series, University Press Group Limited (2006) 28. Reinarz, A., Charrier, D.E., Bader, M., Bovard, L., Dumbser, M., Duru, K., Fambri, F., Gabriel, A.A., Gallard, J.M., Köppel, S., Krenz, L., Rannabauer, L., Rezzolla, L., Samfass, P., Tavelli, M., Weinzierl, T.: ExaHyPE: An engine for parallel dynamically adaptive simulations of wave problems. Computer Physics Communications 254, 107251 (2020). https://doi.org/10.1016/j.cpc.2020.107251 29. Matthew Rocklin: Dask: Parallel Computation with Blocked algorithms and Task Scheduling. In: Kathryn Huff, James Bergstra (eds.) Proceedings of the 14th Python in Science Conference. pp. 126 – 132 (2015). https://doi.org/10.25080/ Majora-7b98e3ed-013 30. Sarri, A., Guillas, S., Dias, F.: Statistical emulation of a tsunami model for sensitivity analysis and uncertainty quantification. Natural Hazards and Earth System Sciences 12(6), 2003–2018 (2012). https://doi.org/10.5194/nhess-12-2003-2012, https:// nhess.copernicus.org/articles/12/2003/2012/ 31. Seelinger, L., Cheng-Seelinger, V., Davis, A., Parno, M., Reinarz, A.: UM-Bridge: Uncertainty quantification and modeling bridge. Journal of Open Source Software 8(83), 4748 (2023). https://doi.org/10.21105/joss.04748, https://doi.org/10.21105/ joss.04748 32. Seelinger, L., Reinarz, A., Lykkegaard, M.B., Akers, R., Alghamdi, A.M., Aristoff, D., Bangerth, W., Bénézech, J., Diez, M., Frey, K., Jakeman, J.D., Jørgensen, J.S., Kim, K.T., Kent, B.M., Martinelli, M., Parno, M., Pellegrini, R., Petra, N., Riis, N.A., Rosenfeld, K., Serani, A., Tamellini, L., Villa, U., Dodwell, T.J., Scheichl, R.: Democratizing uncertainty quantification. Journal of Computational Physics p. 113542 (2024). https://doi.org/https://doi.org/10.1016/j.jcp.2024.113542, https: //www.sciencedirect.com/science/article/pii/S0021999124007903 33. Seelinger, L., Reinarz, A., Rannabauer, L., Bader, M., Bastian, P., Scheichl, R.: High performance uncertainty quantification with parallelized multilevel markov chain monte carlo. In: Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis. SC ’21, Association for Computing Machinery, New York, NY, USA (2021). https://doi.org/10.1145/3458817.3476150, https://doi.org/10.1145/3458817.3476150 34. Sraj, I., Mandli, K.T., Knio, O.M., Dawson, C.N., Hoteit, I.: Uncertainty quantification and inference of manning’s friction coefficients using dart buoy data during

Dynamic Load Balancing for Uncertainty Quantification

19

the tōhoku tsunami. Ocean Modelling 83, 82–97 (2014). https://doi.org/https: //doi.org/10.1016/j.ocemod.2014.09.001, https://www.sciencedirect.com/science/ article/pii/S1463500314001322 35. Stuart, A.M.: Inverse problems: A Bayesian perspective. Acta Numerica 19, 451–559 (2010). https://doi.org/10.1017/S0962492910000061 36. Uphoff, C., Rettenberger, S., Bader, M., Madden, E.H., Ulrich, T., Wollherr, S., Gabriel, A.A.: Extreme scale multi-physics simulations of the tsunamigenic 2004 sumatra megathrust earthquake. In: Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis. SC ’17, Association for Computing Machinery, New York, NY, USA (2017). https://doi. org/10.1145/3126908.3126948, https://doi.org/10.1145/3126908.3126948 37. Weinzierl, T.: The peano software—parallel, automaton-based, dynamically adaptive grid traversals. ACM Trans. Math. Softw. 45(2) (Apr 2019). https://doi.org/10. 1145/3319797, https://doi.org/10.1145/3319797 38. Willoughby, R.A.: Solutions of ill-posed problems (a. n. tikhonov and v. y. arsenin). SIAM Review 21(2), 266–267 (1979). https://doi.org/10.1137/1021044, https:// doi.org/10.1137/1021044