[BUG] - memory leak

[bknight1]

I'm running out of memory when running on the HPC:

s:qurmstepd: error: Detected 1 oom_kill event in StepId=42492733.0. Some of the step tasks have been OOM Killed.
srun: error: nid001653: task 16: Out Of Memory
srun: Terminating StepId=42492733.0
slurmstepd: error: *** STEP 42492733.0 ON nid001653 CANCELLED AT 2026-05-08T22:47:45 ***

I'm running #SBATCH --mem-per-cpu=1200M which is close to the max you can request per cpu.

[lmoresi]

We are going to need a bit more help on the setup How many cpu's Tinestepping ? Particles or not ? L

…

--- 8< ------------------------------ Prof Louis Moresi Computational Geophysics Research School of Earth Sciences The Australian National University +61 2 6125 5590 ***@***.*** TEQSA Provider ID: PRV12002 (Australian University) | CRICOS Provider Code: 00120C The Australian National University acknowledges, celebrates and pays our respects to the Ngunnawal and Ngambri people of the Canberra region and to all First Nations Australians on whose traditional lands we meet and work, and whose cultures are among the oldest continuing cultures in human history.

________________________________ From: Ben Knight ***@***.***> Sent: Saturday, May 9, 2026 1:08:31 PM To: underworldcode/underworld3 ***@***.***> Cc: Subscribed ***@***.***> Subject: [underworldcode/underworld3] [BUG] - memory leak (Issue #176) [https://avatars.githubusercontent.com/u/55677727?s=20&v=4]bknight1 created an issue (underworldcode/underworld3#176)<#176> I'm running out of memory when running on the HPC: s:qurmstepd: error: Detected 1 oom_kill event in StepId=42492733.0. Some of the step tasks have been OOM Killed. srun: error: nid001653: task 16: Out Of Memory srun: Terminating StepId=42492733.0 slurmstepd: error: *** STEP 42492733.0 ON nid001653 CANCELLED AT 2026-05-08T22:47:45 *** I'm running #SBATCH --mem-per-cpu=1200M which is close to the max you can request per cpu. — Reply to this email directly, view it on GitHub<#176>, or unsubscribe<https://github.com/notifications/unsubscribe-auth/ADABPI77H3O55ESPIVBJROL4Z2OK7AVCNFSM6AAAAACYW2JCQ2VHI2DSMVQWIX3LMV43ASLTON2WKOZUGQYTAOBUGQYTQMA>. You are receiving this because you are subscribed to this thread.Message ID: ***@***.***>

[bknight1]

@lmoresi sorry for the vagueness. Here's a basic script showing memory growth that can be used on any number of processors. The usage seems to linearly scale with procs, e.g 1 = 100MB, 2 = 200MB etc

import underworld3 as uw
import numpy as np
import os
import resource
import sympy
import gc
from mpi4py import MPI

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

def get_memory_usage():
    try:
        if os.uname().sysname == 'Darwin':
            rss = int(os.popen('ps -p %d -o rss=' % os.getpid()).read())
            return rss / 1024.0 # MiB
        else:
            with open('/proc/self/status') as f:
                for line in f:
                    if line.startswith('VmRSS:'):
                        return int(line.split()[1]) / 1024.0 # MiB
    except:
        return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024.0

# Parameters
res = 32
n_steps = 50

# Setup mesh
mesh = uw.meshing.StructuredQuadBox(elementRes=(res, res))
mesh._dminterpolation_cache.max_entries = 10 # CAP CACHE

# Variables
v = uw.discretisation.MeshVariable("u", mesh, mesh.dim, degree=2)
p = uw.discretisation.MeshVariable("p", mesh, 1, degree=1)
T = uw.discretisation.MeshVariable("T", mesh, 1, degree=1)

# Swarm for advection
swarm = uw.swarm.Swarm(mesh)
swarm.populate(fill_param=2)

# Stokes
stokes = uw.systems.Stokes(mesh, velocityField=v, pressureField=p)
stokes.constitutive_model = uw.constitutive_models.ViscousFlowModel
stokes.constitutive_model.Parameters.viscosity = 1.0
stokes.bodyforce = sympy.Matrix([0, T.sym[0]])
stokes.add_dirichlet_bc([0.0, 0.0], "Top")
stokes.add_dirichlet_bc([0.0, 0.0], "Bottom")
stokes.add_dirichlet_bc([0.0, sympy.oo], "Left")
stokes.add_dirichlet_bc([0.0, sympy.oo], "Right")

# AdvDiff
advdiff = uw.systems.AdvDiffusion(mesh, u_Field=T, V_fn=v)
advdiff.constitutive_model = uw.constitutive_models.DiffusionModel
advdiff.constitutive_model.Parameters.diffusivity = 1.0
advdiff.add_dirichlet_bc([0.0], "Top")
advdiff.add_dirichlet_bc([1.0], "Bottom")

# Initial T
T.data[:] = 0.0
with mesh.access(T):
    T.data[:, 0] = 1.0 - mesh.data[:, 1]

if rank == 0:
    print(f"Starting Stokes + AdvDiff leak test for {n_steps} steps on {size} ranks...")
    print(f"Resolution: {res}x{res}")
    print(f"Initial Memory = {get_memory_usage():.2f} MiB")

mem_history = []

for i in range(n_steps):
    # 1. Solve Stokes
    # stokes.solve()

    # dt = stokes.estimate_dt()
    dt = 0.001

    # 2. Solve AdvDiff
    advdiff.solve(timestep=dt)

    # 3. Particle advection (triggers interpolation and coordinate caching)
    # We'll use a dummy advection to change coordinates
    with swarm.access(swarm):
        swarm.data[...] = swarm.data[...] + 0.01 * np.random.rand(*swarm.data.shape)
        swarm.data[...] = np.clip(swarm.data[...], 0, 1)
    
    # Evaluate velocity at swarm to trigger cache population
    v_at_swarm = uw.function.evaluate(v.sym, swarm.data)

    gc.collect()

    if i % 5 == 0:
        local_mem = get_memory_usage()
        total_mem = comm.reduce(local_mem, op=MPI.SUM, root=0)

        if rank == 0:
            dm_fields = mesh.dm.getNumFields()
            print(f"Step {i:3d}: Total Memory = {total_mem:7.2f} MiB | DM Fields = {dm_fields:2d}")
            mem_history.append(total_mem)

if rank == 0:
    print("\nMemory usage summary (Total across ranks):")
    for i, mem in enumerate(mem_history):
        print(f"Step {i*5:3d}: {mem:7.2f} MiB")

    growth = mem_history[-1] - mem_history[0]
    print(f"\nTotal growth: {growth:.2f} MiB")
    if growth > 5.0:
        print("\nRESULT: Memory leak detected!")
    else:
        print("\nRESULT: Memory usage stable.")

[lmoresi]

@bknight1 — we can't fix the leak from your initial report alone, but PR #179 adds uw.utilities.memprobe, a small diagnostic module that should help us pin which subsystem is growing on your HPC run.

Quickest way to use it:

UW_MEMPROBE=1 mpirun -n 16 python your_script.py

That turns on per-Stokes.solve() and per-NavierStokes.solve() diffs that look like:

[memprobe] Stokes.solve:
  RSS +0.42 MiB
  kdtree: live +1, total_constructed +1

A clean run shows mostly zero deltas. A leak shows the same counter growing on every step. For a more detailed view, wrap your step loop:

from underworld3.utilities import memprobe
with memprobe.probe(f"step {step}", full=True):
    ...  # your step body

full=True walks gc.get_objects() and shows Python-class-level growth sorted by absolute change — usually the dominant suspect surfaces immediately. If RSS grows but no Python class does, the leak is in PETSc/C territory and the guide docs/developer/guides/memory-diagnostics.md walks through -log_view -malloc_dump.

Could you give it a run on the workload that OOM'd and paste back the per-step output? Even ~20 steps' worth would be enough to spot a linear growth signature.

[bknight1]

The changes in #178 managed to approximately double the number of timesteps, but still running into a OOM error. Will use the memprobe module to diagnose further

[bknight1]

Here's some memprobe info:

[memprobe] step 0:
  RSS +63.29 MiB
  kdtree: live +1, total_constructed +102
  py_classes (top 10 by |Δ|):
    underworld3.utilities.nd_array_callback.NDArray_With_Callback: +19
    underworld3.function.expressions.UWexpression: +12
    underworld3.swarm.SwarmVariable: +6
    underworld3.utilities._jitextension.Functions: +6
    underworld3.function._dminterp_wrapper.Enum: +5
    underworld3.function._dminterp_wrapper.CachedDMInterpolationInfo: +5
    underworld3.discretisation.discretisation_mesh_variables.MeshVariable_ij: +2
    underworld3.discretisation.discretisation_mesh_variables._BaseMeshVariable: +1
    underworld3.ckdtree.memoryview: +1
    underworld3.utilities._jitextension._JITConstant: -1

[memprobe] step 1:
  RSS +8.55 MiB
  kdtree: total_constructed +82
  py_classes (top 10 by |Δ|):
    underworld3.function._dminterp_wrapper.CachedDMInterpolationInfo: +5
    underworld3.utilities.nd_array_callback.NDArray_With_Callback: -3
    underworld3.swarm.SwarmVariable: -3

[memprobe] step 2:
  RSS +6.00 MiB
  kdtree: total_constructed +82
  py_classes (top 10 by |Δ|):
    underworld3.swarm.SwarmVariable: -4
    underworld3.utilities.nd_array_callback.NDArray_With_Callback: -4

[memprobe] step 3:
  RSS +2.00 MiB
  kdtree: total_constructed +82
  py_classes (top 10 by |Δ|):
    underworld3.swarm.SwarmVariable: +9
    underworld3.utilities.nd_array_callback.NDArray_With_Callback: +9

[memprobe] step 4:
  RSS +2.00 MiB
  kdtree: total_constructed +82
  py_classes (top 10 by |Δ|):
    underworld3.swarm.SwarmVariable: +11
    underworld3.utilities.nd_array_callback.NDArray_With_Callback: +11

[memprobe] step 5:
  RSS +8.02 MiB
  kdtree: total_constructed +82
  py_classes (top 10 by |Δ|):
    underworld3.swarm.SwarmVariable: -9
    underworld3.utilities.nd_array_callback.NDArray_With_Callback: -9

[memprobe] step 6:
  RSS +4.20 MiB
  kdtree: total_constructed +82
  py_classes (top 10 by |Δ|):
    underworld3.swarm.SwarmVariable: -4
    underworld3.utilities.nd_array_callback.NDArray_With_Callback: -4

[lmoresi]

Thanks @bknight1 — the trace is useful. Three observations:

1. KDTree construction count: clean cross-validation for #182

kdtree: total_constructed +82  (per step)

After step 0 the count climbs by ~82 every step, while live stays bounded. That's exactly the "rebuild-and-discard" pattern your #182 (KDTree consolidation via _get_kdtree()) is targeting. Once #182 lands and the kdtree cache is wired in, this line should drop to near-zero on a warm loop. Good signal that #182 is solving a real, measurable problem.

2. Residual ~5 MiB/step is PETSc-managed memory

RSS still climbs ~5 MiB/step from step 3 onward, but the Python-class deltas oscillate around zero (SwarmVariable + NDArray_With_Callback go +9/-9 etc.). No Python class is growing linearly. That means the residual leak is not Python-object accumulation — it's in PETSc-managed memory (Mat/Vec/DM internals, compiled-context cache, etc.) or in C extensions that don't expose Python-level counters.

The Python-side tracker is doing its job here by ruling itself out. Next step is -log_view -malloc_dump. Easiest from your script:

from underworld3.utilities import memprobe
memprobe.dump_petsc_leaks_at_finalize()   # call once, near the top

Then run a short (e.g. 20-step) job — PETSc will dump unfreed allocations at finalize. The output is verbose but the offenders will repeat across runs and stand out.

3. Step-0 anomaly is expected

Step 0 shows +63 MiB RSS and +102 constructed, vs. ~82 per subsequent step. That's first-touch / warm-up — JIT compilation, DM build, KDtree caches. Just confirm your test/assertion budgets are measured on post-warmup steps (step 1+).

[bknight1]

I've just done the petsc leak, here's the summary:

The following objects were never freed
-----------------------------------------
[0] PetscSpace sum P2
[0] DM shell DM_0x84000002_477
[0] PetscSF (null) PetscSF_0x84000002_478
[0] PetscDS (null) PetscDS_0x84000002_479
[0] PetscWeakForm (null) PetscWeakForm_0x84000002_480
[0] PetscSpace poly P2
[0] DM plex uw_hierarchical_dm
[0] PetscSF basic PetscSF_0x84000003_214