Reuse linear rotated free-slip solver workspace - #418
Conversation
Fix repeated transient Stokes solves rebuilding the rotated PtAP operator, fieldsplit Schur preconditioner, and GAMG hierarchy on every timestep. Cache the linearity classification, rotation matrices, transformed operator, pressure mass, nullspaces, KSP/PC, and deterministic operator probes. Skip Jacobian assembly for RHS-only updates, refresh safely when constitutive coefficient MeshVariable state or explicit physical time changes, and destroy all owned PETSc state during solver rebuilds. Store the Cartesian reaction directly for rotated topography instead of retaining native operator/RHS copies. Expose the existing physical time argument through the high-level Stokes.solve wrapper. Add regression coverage for RHS reuse, viscosity-driven invalidation, explicit-time refresh, and inverse-viscosity scaling. Document lifecycle, invalidation, MPI launcher requirements, and validation for issue underworldcode#417.
Bounded production A1 restart: timing and memoryI reran the actual Zhong A1 production path after this fix, rather than only the synthetic repeated-solve probe. Configuration:
The run completed normally in 87.79 seconds and neither guard was reached. The external RSS value is the first sample after each Stokes completion.
Timing summary
Interpretation
Step-510 diagnostics continue the checkpoint trajectory: paper time 0.242451, RMS velocity 57.8450, surface/CMB Nu 2.69543/3.20704, and mean temperature 0.317257. This validates the fix over a bounded ten-step production continuation. It does not yet prove long-run whole-process stability; a guarded run of at least 100 steps on a larger-memory host remains appropriate. Full benchmark report: results_bench_100_zhong2008_thermal_convection_A1.py.md |
Guarded production follow-up: 100 additional A1 stepsI extended the real Zhong A1
Each process includes one extra restart-state Stokes solve, excluded from the advancing-step timing. In the first process, aggregate RSS reached 3220.98 MiB after step 559 and then released about 644 MiB at step 560. The second process released about 141 MiB at step 570; its external RSS slope over steps 590-609 was 0.74 MiB/step aggregate. Warm Stokes-local deltas were mostly zero or sub-MiB and included periodic net-negative releases. This is unlike the former per-step The physical trajectory remained continuous. At step 610 (
This is deliberately a segmented 100-step validation, not a single uninterrupted 100-step memory profile. It supports the cache/reuse fix and validates a guarded checkpoint workflow; it does not justify an unmonitored full run on this machine. Full logs, RSS interpretation, and checkpoint metrics are documented in benchmark commit |
Extended production validation: checkpoint 810The real Zhong A1
The four processes completed in 1172.39 s (19.54 min). Advancing solve intervals averaged 5.37 s, with a 5.21 s median. The largest process peak was 3.74 GiB, less than half the configured limit. Warm Stokes-local RSS repeatedly flattened to zero/sub-MiB changes and included multiple collective releases rather than cumulative workspace growth. Physical checkpoint results remained continuous:
Guarded post-fix production evidence now covers 310 evolved steps after checkpoint 500. Full details are in benchmark commit |
Adversarial reviewStrong PR: the ownership discipline is careful (explicit destroys, Probe results1. Rampable UWexpression-constant coefficients — FAIL, blocking (silent wrong answer)Viscosity expressed purely as a rampable constant (
Fix shape: include the float values of rampable atomic constants reachable from the constitutive parameters / penalty / constraints in the invalidation key alongside the 2.
|
… with a constants-aware key (#417, supersedes #418) (#543) * Reuse the rotated free-slip solver workspace across solves (#417) Repeated transient rotated free-slip Stokes solves rebuilt the rotation Q, the PtAP'd operator, the fieldsplit Schur KSP/PC and the GAMG hierarchy on every timestep — an allocator high-water problem (RSS growth to OOM, issue #417) and roughly half the per-step cost of the Zhong A1 production runs. Cache the workspace on the solver between solves, re-derived from the original PR #418 (bec76bb) at the seam of the rewritten unified Newton loop: - geometry tier (Q/Qt, constrained rows, fault pair blocks, custom-FMG prolongation) reused while the boundary/fault registration and DM are unchanged; - structure tier (Ahat, Schur pmat, nullspaces, KSP/PC) reused as objects with values refreshed in place — the loop's own between-iteration pattern; - an iteration-0 fast path skips Jacobian assembly / ptap / PCSetUp entirely when the operator coefficient state counters match (RHS-only timesteps); - the cache is forfeited for direct-LU, prescribed-datum and fault interface-law solves, and torn down by _reset() and the _build() full rebuild before the SNES/DM are destroyed; - expose the existing time= argument through the Stokes.solve wrapper; an explicit time vetoes the fast path (petsc_t bypasses every counter). Unlike the original one-shot linear path, a wrong fast-path verdict here cannot return a stale solution: the loop measures the true residual at every iterate and reassembles from iteration 1 on. Regression: RHS-only reuse (same Q/Ahat/KSP handles), viscosity-field invalidation with in-place refresh, and the time= veto, in test_1018_rotated_freeslip. Production evidence on the original mechanism: 310 guarded Zhong A1 steps, flat RSS (PR #418 thread). Underworld development team with AI support from Claude Code * Close the rotated-workspace blind spot: rampable constants join the key, the verdict goes collective The PR #418 review's unresolved finding: the reuse key was built from MeshVariable._state counters, which are BLIND to rampable UWexpression constants — the #416 contract lets a constant change value with no state bump, so a 2x viscosity ramp between solves reported "unchanged" and (on the original one-shot path) returned a bit-identical stale solution while the matrix-probe safety net was disabled by the very verdict it was meant to check. The reworked loop already made a stale verdict non-fatal (the true residual is measured every iterate and iteration 1 onward always reassembles), but the verdict itself must still be honest: - the operator key now includes the packed constants[] values the kernels will actually assemble with, plus the JIT bundle key (covers an in-place kernel rewire). Measured on the ramp probe: the naive key reported workspace_reused=True for a 2x constant ramp; this key reports False and reassembles. Over-invalidation on RHS-only constant changes is accepted — reassembly is the safe default; - if the constants manifest or the coefficient enumeration cannot be read, the fast path is forfeited outright — correctness first; - the fast path additionally requires a self-measured linear hint (last solve converged in <= 1 increment): for a nonlinear model the cached operator is the previous solve's tangent, and the skip would only trade an assembly for a wasted increment; - the match verdict is allgathered and must be unanimous before it gates any collective PETSc call — state counters follow rank-local writes, and a rank-divergent verdict is a deadlock; - on a detected change the stored signature is poisoned before the in-place refresh, so an exception mid-refresh cannot leave a stale key that later matches half-updated values. Regression: test_rotated_workspace_constant_ramp_invalidates (fail-before validated on the naive key: the reused flag lies there) with its own armed- fast-path negative control and a fresh-solver control at the ramped viscosity; test_rotated_workspace_deform_invalidates re-proves the mesh.deform teardown on the reworked cache. Underworld development team with AI support from Claude Code * Fault contact composes with the rotated workspace cache: pair blocks cache, interface laws opt out solve_with_fault drives the same rotated Newton loop, so the cross-solve workspace decision had to be made explicitly for the fault machinery: - FRICTIONLESS pair blocks are geometry (coincident-node pairing + fault normals live in Q, keyed by the fault registration in the geometry signature) — they cache. A warm repeat reuses the rotation; a cold re-solve rides the iteration-0 fast path; both match a fresh-solver control on the same mesh. - INTERFACE-LAW solvers (viscous / Coulomb / rate-state) opt out entirely: the interface tangent is reassembled per iterate at the current slip rates and the reaction-fed normal stress is Picard-lagged solver state — neither is keyable registration state, so cache_allowed excludes them at the top of solve_rotated_freeslip. Regression: test_fault_repeat_solve_composes_with_workspace_reuse covers both arms, with a fresh-solver control for the cached arm and the absent- cache assertion for the opt-out. Underworld development team with AI support from Claude Code * Review response: persist the operator key only for values Ahat holds; null before destroy Two minors from the #543 review. (m2) The persisted operator_sig now requires that this solve either reassembled Ahat or rode a key-matched fast path - a poisoned solve exiting at iteration 0 without assembling can no longer store a fresh key against unrefreshed values (the window was bounded by the always-reassemble net; now it is closed). (m4) Both teardown sites null the cache attribute before destroying, so an exception mid-destroy leaves objects unreachable rather than arming a double-destroy. Underworld development team with AI support from Claude Code --------- Co-authored-by: Tyagi <gthyagi@gmail.com>
|
Superseded by #543 (merged as e475246), which carries your commit with authorship preserved plus the rework onto the current rotated loop and the constants-aware invalidation key from the July review. Your 310-step Zhong A1 production validation remains the load-bearing evidence for the mechanism and is credited in the PR — when you have a slot, a 50-step guarded A1 restart segment on current development would re-anchor it against the reworked loop. Thanks for the mechanism and the exceptional validation trail. |
2 participants
Summary
Fixes the unbounded resident-memory growth observed when a linear rotated free-slip Stokes system is solved repeatedly, as in transient mantle-convection models.
The change caches the rotated linear solver workspace and reuses it while the operator is unchanged. It also adds explicit invalidation for coefficient, time, boundary-condition, constitutive-model, solver, and DM changes.
Tracks #417.
Root cause
Repeated solves rebuilt two expensive classes of PETSc state:
_residual_is_nonlinear()assembled and compared two Jacobians on every solve.Stage-isolation measurements showed approximately 76.77 MiB growth across 12 repeated nonlinearity probes and 117.23 MiB across 12 direct rotated solves. Explicitly destroying only the previous result vectors did not address the retained PETSc solver state.
Implementation
Q,Q^T, transformed operators, pressure mass, nullspaces, KSP/PC, coefficient state, and deterministic operator probes.MeshVariable._statevalues are unchanged.time=is supplied, coveringmesh.t-dependent symbolic operators.time=separately fromtimestep=inStokes.solve().The design and invalidation contract are documented in
docs/developer/design/ROTATED_FREESLIP_LINEAR_REUSE.md.Validation
The isolated fix was cherry-picked onto the current
developmenttip. The branch is exactly one commit ahead with no merge conflicts;git diff --checkand Python syntax checks pass.Focused serial and MPI validation covered repeated RHS solves, viscosity invalidation, explicit-time refresh, spherical velocity and leakage, Zhong 3-D reaction topography, boundary traction, dynamic topography, and partition independence.
Before this change, the serial spherical probe grew by 102.03 MiB over 12 solves (about 8.54 MiB/solve). The production 8-rank Zhong A1 run had reached 54.29-59.18 GiB aggregate RSS on a 16 GiB machine and triggered macOS watchdog restarts.
MPI checks used the Pixi environment's Open MPI launcher. The Homebrew Open MPI launcher is ABI/version-mismatched with the environment and can hang during startup; no
FI_PROVIDER=tcpworkaround was used.Scope
This PR changes only the UW3 rotated free-slip solve lifecycle, focused tests, and developer documentation. Benchmark-repository scripts and result documents are not included.