From e093e316f4e952ec567042167c9fe434a1ab626a Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Sat, 22 Aug 2026 11:55:30 +0200 Subject: [PATCH 1/3] Added weak scaling --- .../cube_weak_scaling/params_poisson.py | 282 ++++++++++++++++++ profiling/submit_poisson_weak_scaling.py | 52 ++++ 2 files changed, 334 insertions(+) create mode 100644 profiling/examples/Poisson/cube_weak_scaling/params_poisson.py create mode 100644 profiling/submit_poisson_weak_scaling.py diff --git a/profiling/examples/Poisson/cube_weak_scaling/params_poisson.py b/profiling/examples/Poisson/cube_weak_scaling/params_poisson.py new file mode 100644 index 000000000..96482de1b --- /dev/null +++ b/profiling/examples/Poisson/cube_weak_scaling/params_poisson.py @@ -0,0 +1,282 @@ +# ----------------------------- +# Description of the simulation +# ----------------------------- +# Please fill in a verbal description of the simulation. +# It will be printed at the beginning of the simulation and can be used to keep track of the different runs. + +name = "Poisson weak scaling on 3D cuboid" +description = """ +Weak scaling test for Poisson equation on a 3D cuboid. +The manufactured solution is a simple product of sines and cosines. +Homogeneous Dirichlet boundary conditions are set in direction x. +""" + +import logging + +from struphy import set_logging_level + +set_logging_level(logging.WARNING) + +import argparse + +from mpi4py import MPI + +# ------------------ +# Import Struphy API +# ------------------ +from struphy import ( + BaseUnits, + DerhamOptions, + EnvironmentOptions, + Simulation, + Time, + domains, + grids, +) + +# --------------------- +# Instance of the model +# --------------------- +from struphy.models import Poisson + +# Units +base_units = BaseUnits() + +# Model instance +model = Poisson(base_units=base_units) + +# -------------------------- +# Instance of the simulation +# -------------------------- + +# Environment options +# `--id` distinguishes runs that share a rank count but differ in something else; the +# profiling driver passes its launch counter (see `ProfilingJob.build_commands`). +# Unknown flags are ignored so the driver can forward other parameters as well. +parser = argparse.ArgumentParser() +parser.add_argument("--id", type=int, default=0, help="Run id, used to name the output folder.") +parser.add_argument( + "--cells-per-rank", + type=int, + default=64, + help="Number of cells assigned to each rank in the weak-scaling direction.", +) +parser.add_argument( + "--verify", + action="store_true", + help="Save fields, post-process them, and check the manufactured solution.", +) +args, _ = parser.parse_known_args() + +mpi_size = MPI.COMM_WORLD.Get_size() + +# List all variables and decide whether to save their data. +# Full field output is only needed for the small verification mode; production weak +# scaling runs keep output compact and rely on profiling_data.h5. +model.em_fields.phi.save_data = args.verify +model.em_fields.source.save_data = args.verify + +env = EnvironmentOptions( + sim_folder=f"sim_{args.id:02d}", + profiling_activated=True, + restart=False, +) + +# Time stepping +time_opts = Time() + +# Geometry +Lx = 2.0 * mpi_size +Ly = 3.0 +Lz = 4.0 +domain = domains.Cuboid(r1=Lx, l2=-Ly/2, r2=Ly/2, r3=Lz) + +# Fluid equilibrium (can be used as part of initial conditions) +equil = None + +# Grid +num_elements = (args.cells_per_rank * mpi_size, args.cells_per_rank, args.cells_per_rank) +grid = grids.TensorProductGrid(num_elements=num_elements, mpi_dims_mask=(True, False, False)) + +# Derham options +derham_opts = DerhamOptions(degree=(1, 2, 3), bcs=(("dirichlet", "dirichlet"), None, None)) + +# Simulation object +sim = Simulation( + model=model, + name=name, + description=description, + params_path=__file__, + env=env, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, +) + +# ------------------ +# Propagator options +# ------------------ + +from struphy.linear_algebra.solver import SolverParameters +solver_params = SolverParameters(tol=1e-8, maxiter=3000, info=True, recycle=True) +model.propagators.poisson.options = model.propagators.poisson.Options(stab_eps=0.0, + solver="pcg", + precond=None, + solver_params=solver_params, + ) + +# ------------------ +# Initial conditions +# ------------------ +import numpy as np +from struphy.initial.base import GenericPerturbation + +def exact_solution(x, y, z): + return np.sin(np.pi / Lx * x) * np.cos(2 * np.pi / Ly * y + 2 * np.pi / Lz * z) + +def rhs_fun(x, y, z): + return exact_solution(x, y, z) * ((np.pi / Lx) ** 2 + (2 * np.pi / Ly) ** 2 + (2 * np.pi / Lz) ** 2) + +rhs_perturbation = GenericPerturbation(rhs_fun, given_in_basis="physical") + +model.em_fields.source.add_perturbation(rhs_perturbation) + + +if __name__ == "__main__": + sim.run(one_time_step=True) + + if not args.verify: + if sim.comm.rank == 0: + import os + + results_dir = os.path.join(sim.env.path_out, "results") + os.makedirs(results_dir, exist_ok=True) + + np.save(os.path.join(results_dir, "resolution.npy"), sim.grid.num_elements) + np.save(os.path.join(results_dir, "spline_degree.npy"), sim.derham_opts.degree) + np.save(os.path.join(results_dir, "mpi_size.npy"), mpi_size) + np.save(os.path.join(results_dir, "cells_per_rank.npy"), args.cells_per_rank) + raise SystemExit(0) + + sim.pproc(parallel_pproc=True) + + def plot_slices(num, exact, name, slice_pt_x=0, slice_pt_y=0, slice_pt_z=0): + from matplotlib import pyplot as plt + + fig = plt.figure(figsize=(16, 12)) + + plt.subplot(3, 3, 1) + plt.pcolor(y[slice_pt_x, :, :], z[slice_pt_x, :, :], num[slice_pt_x, :, :]) + plt.colorbar() + plt.xlabel("y") + plt.ylabel("z") + plt.title("{} from struphy, slice at x = {:.2f}".format(name, x[slice_pt_x, 0, 0])) + + plt.subplot(3, 3, 4) + plt.pcolor(y[slice_pt_x, :, :], z[slice_pt_x, :, :], exact(x[slice_pt_x, :, :], y[slice_pt_x, :, :], z[slice_pt_x, :, :])) + plt.colorbar() + plt.xlabel("y") + plt.ylabel("z") + plt.title("{} exact, slice at x = {:.2f}".format(name, x[slice_pt_x, 0, 0])) + + plt.subplot(3, 3, 7) + plt.pcolor(y[slice_pt_x, :, :], z[slice_pt_x, :, :], np.abs(num[slice_pt_x, :, :] - exact(x[slice_pt_x, :, :], y[slice_pt_x, :, :], z[slice_pt_x, :, :]))) + plt.colorbar() + plt.xlabel("y") + plt.ylabel("z") + plt.title("{} error, slice at x = {:.2f}".format(name, x[slice_pt_x, 0, 0])) + + plt.subplot(3, 3, 2) + plt.pcolor(x[:, slice_pt_y, :], z[:, slice_pt_y, :], num[:, slice_pt_y, :]) + plt.colorbar() + plt.xlabel("x") + plt.ylabel("z") + plt.title("{} from struphy, slice at y = {:.2f}".format(name, y[0, slice_pt_y, 0])) + + plt.subplot(3, 3, 5) + plt.pcolor(x[:, slice_pt_y, :], z[:, slice_pt_y, :], exact(x[:, slice_pt_y, :], y[:, slice_pt_y, :], z[:, slice_pt_y, :])) + plt.colorbar() + plt.xlabel("x") + plt.ylabel("z") + plt.title("{} exact, slice at y = {:.2f}".format(name, y[0, slice_pt_y, 0])) + + plt.subplot(3, 3, 8) + plt.pcolor(x[:, slice_pt_y, :], z[:, slice_pt_y, :], np.abs(num[:, slice_pt_y, :] - exact(x[:, slice_pt_y, :], y[:, slice_pt_y, :], z[:, slice_pt_y, :]))) + plt.colorbar() + plt.xlabel("x") + plt.ylabel("z") + plt.title("{} error, slice at y = {:.2f}".format(name, y[0, slice_pt_y, 0])) + + plt.subplot(3, 3, 3) + plt.pcolor(x[:, :, slice_pt_z], y[:, :, slice_pt_z], num[:, :, slice_pt_z]) + plt.colorbar() + plt.xlabel("x") + plt.ylabel("y") + plt.title("{} from struphy, slice at z = {:.2f}".format(name, z[0, 0, slice_pt_z])) + + plt.subplot(3, 3, 6) + plt.pcolor(x[:, :, slice_pt_z], y[:, :, slice_pt_z], exact(x[:, :, slice_pt_z], y[:, :, slice_pt_z], z[:, :, slice_pt_z])) + plt.colorbar() + plt.xlabel("x") + plt.ylabel("y") + plt.title("{} exact, slice at z = {:.2f}".format(name, z[0, 0, slice_pt_z])) + + plt.subplot(3, 3, 9) + plt.pcolor(x[:, :, slice_pt_z], y[:, :, slice_pt_z], np.abs(num[:, :, slice_pt_z] - exact(x[:, :, slice_pt_z], y[:, :, slice_pt_z], z[:, :, slice_pt_z]))) + plt.colorbar() + plt.xlabel("x") + plt.ylabel("y") + plt.title("{} error, slice at z = {:.2f}".format(name, z[0, 0, slice_pt_z])) + + return fig + + if sim.comm.rank == 0: + sim.load_plotting_data() + + Tstart = sim.t_grid[0] + rhs_data = sim.spline_values.em_fields.source_log + print(rhs_data) + rhs = rhs_data.data[Tstart][0] + + Tend = sim.t_grid[-1] + phi_data = sim.spline_values.em_fields.phi_log + print(phi_data) + phi = phi_data.data[Tend][0] + x = sim.grids_phy[0] + y = sim.grids_phy[1] + z = sim.grids_phy[2] + + slice_pt_x = x.shape[0] // 2 + slice_pt_y = y.shape[1] // 2 + slice_pt_z = 0 + + fig_rhs = plot_slices(rhs, rhs_fun, "RHS", slice_pt_x=slice_pt_x, slice_pt_y=slice_pt_y, slice_pt_z=slice_pt_z) + fig_phi = plot_slices(phi, exact_solution, "Phi", slice_pt_x=slice_pt_x, slice_pt_y=slice_pt_y, slice_pt_z=slice_pt_z) + + rel_err_rhs = np.max(np.abs(rhs - rhs_fun(x, y, z))) / np.max(np.abs(rhs_fun(x, y, z))) + rel_err_phi = np.max(np.abs(phi - exact_solution(x, y, z))) / np.max(np.abs(exact_solution(x, y, z))) + + print(f"Max relative error in RHS: {rel_err_rhs:.2e}") + print(f"Max relative error in Phi: {rel_err_phi:.2e}") + + assert rel_err_rhs < 5e-3, f"The computed RHS does not match the exact RHS, max rel error = {rel_err_rhs}." + assert rel_err_phi < 5e-3, f"The computed solution does not match the exact solution, max rel error = {rel_err_phi}." + + import os + # `path_out` is the run's output folder; `sim_folder` alone is a bare name + # resolved against the CWD. The profiling packaging picks these files up from + # here and uploads them as `results-run`. + results_dir = os.path.join(sim.env.path_out, "results") + os.makedirs(results_dir, exist_ok=True) + + np.save(os.path.join(results_dir, "rel_err_rhs.npy"), rel_err_rhs) + np.save(os.path.join(results_dir, "rel_err_phi.npy"), rel_err_phi) + np.save(os.path.join(results_dir, "resolution.npy"), sim.grid.num_elements) + np.save(os.path.join(results_dir, "spline_degree.npy"), sim.derham_opts.degree) + np.save(os.path.join(results_dir, "mpi_size.npy"), mpi_size) + np.save(os.path.join(results_dir, "cells_per_rank.npy"), args.cells_per_rank) + + fig_rhs.savefig(os.path.join(results_dir, "rhs_slices.png")) + fig_phi.savefig(os.path.join(results_dir, "phi_slices.png")) diff --git a/profiling/submit_poisson_weak_scaling.py b/profiling/submit_poisson_weak_scaling.py new file mode 100644 index 000000000..e0bb57f96 --- /dev/null +++ b/profiling/submit_poisson_weak_scaling.py @@ -0,0 +1,52 @@ +import argparse +from pathlib import Path + +from profiling_job import ProfilingCase + + +def main() -> None: + + parser = argparse.ArgumentParser( + description=("Submit Poisson weak scaling profiling jobs and package the results."), + ) + parser.add_argument( + "--upload", + action="store_true", + help="Upload the packaged profiling results to the profiling-data repo.", + ) + parser.add_argument( + "--cells-per-rank", + type=int, + default=64, + help="Number of cells assigned to each rank in the weak-scaling direction.", + ) + args = parser.parse_args() + + script_dir = Path(__file__).resolve().parent + params_dir = script_dir / "examples" / "Poisson" / "cube_weak_scaling" + + profiling_case = ProfilingCase( + label="poisson_cube_weak_scaling", + name="Poisson on cuboid weak scaling test", + description=( + "Weak scaling of the Poisson model with a manufactured solution on a " + "3D cuboid. The grid grows in x with the MPI rank count." + ), + physics_problem="Occurs in many plasma applications.", + struphy_model_used="Poisson", + params_source=params_dir / "params_poisson.py", + language="fortran", + compiler="GNU", + upload=args.upload, + ) + + param_flags = ["--cells-per-rank", str(args.cells_per_rank)] + + for num_tasks in (1, 2, 4, 8, 16, 32, 64, 128, 256): + profiling_case.launch(num_tasks, param_flags=param_flags) + + profiling_case.finalize_run() + + +if __name__ == "__main__": + main() From 0c1cfed5b4a8df07d5e5ac5364f725ff4724e8f8 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Sat, 22 Aug 2026 16:18:26 +0200 Subject: [PATCH 2/3] Fix sim.run(...) --- profiling/examples/Poisson/cube_weak_scaling/params_poisson.py | 3 +-- profiling/submit_poisson_weak_scaling.py | 2 +- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/profiling/examples/Poisson/cube_weak_scaling/params_poisson.py b/profiling/examples/Poisson/cube_weak_scaling/params_poisson.py index 96482de1b..ec6499422 100644 --- a/profiling/examples/Poisson/cube_weak_scaling/params_poisson.py +++ b/profiling/examples/Poisson/cube_weak_scaling/params_poisson.py @@ -78,7 +78,6 @@ env = EnvironmentOptions( sim_folder=f"sim_{args.id:02d}", - profiling_activated=True, restart=False, ) @@ -145,7 +144,7 @@ def rhs_fun(x, y, z): if __name__ == "__main__": - sim.run(one_time_step=True) + sim.run(profiling_activated=True, one_time_step=True) if not args.verify: if sim.comm.rank == 0: diff --git a/profiling/submit_poisson_weak_scaling.py b/profiling/submit_poisson_weak_scaling.py index e0bb57f96..f17ff7f4e 100644 --- a/profiling/submit_poisson_weak_scaling.py +++ b/profiling/submit_poisson_weak_scaling.py @@ -42,7 +42,7 @@ def main() -> None: param_flags = ["--cells-per-rank", str(args.cells_per_rank)] - for num_tasks in (1, 2, 4, 8, 16, 32, 64, 128, 256): + for num_tasks in (1, 2, 4): #, 8, 16, 32, 64, 128, 256): profiling_case.launch(num_tasks, param_flags=param_flags) profiling_case.finalize_run() From a415ff79cca52147834bb428f4f44ba906a1173f Mon Sep 17 00:00:00 2001 From: maxlin Date: Sat, 22 Aug 2026 17:05:36 +0200 Subject: [PATCH 3/3] Add setup for raven --- profiling/clusters.py | 15 +++++++++++++++ profiling/profiling_job.py | 15 +++++++++++++-- profiling/submit_poisson_weak_scaling.py | 9 ++++++++- setup/modules.sh | 7 +++++++ 4 files changed, 43 insertions(+), 3 deletions(-) diff --git a/profiling/clusters.py b/profiling/clusters.py index d64c9a1f6..55a6b6aa3 100644 --- a/profiling/clusters.py +++ b/profiling/clusters.py @@ -57,6 +57,21 @@ def detect_machine_name() -> str | None: SLURM_PRESETS: dict[str, dict] = { + "raven": { + "cpus_per_task": 1, + # Raven's submission filter requires total memory (not --mem-per-cpu) + # for shared jobs. ProfilingCase converts this per-rank value to --mem. + "mem_per_rank_gb": 3, + "max_mem_per_node_gb": 110, + "partition": "general", + "qos": "debug", + "account": "ipp_cpu", + "chdir": "./", + "output": "./%x.%j.out", + "error": "./%x.%j.err", + "mail_type": "none", + "time": "00:15:00", + }, "pitagora_dcgp": { # "nodes": 1, # Should be set by ProfilingCase.launch() # "ntasks_per_node": 1, # Should be set by ProfilingCase.launch() diff --git a/profiling/profiling_job.py b/profiling/profiling_job.py index 8b9b97f5e..62d99d8f4 100644 --- a/profiling/profiling_job.py +++ b/profiling/profiling_job.py @@ -252,9 +252,20 @@ def launch( # Pick the cluster preset for this run, either from the caller's override or the default. if slurm_presets is not None: - cluster_preset = slurm_presets[cluster_name] + cluster_preset = dict(slurm_presets[cluster_name]) else: - cluster_preset = SLURM_PRESETS[cluster_name] + cluster_preset = dict(SLURM_PRESETS[cluster_name]) + + # Some schedulers (notably Raven's submission filter) reject + # --mem-per-cpu for shared jobs. Let their preset express the same + # scaling policy while emitting the required total --mem value. + mem_per_rank_gb = cluster_preset.pop("mem_per_rank_gb", None) + max_mem_per_node_gb = cluster_preset.pop("max_mem_per_node_gb", None) + if mem_per_rank_gb is not None: + memory_gb = mem_per_rank_gb * (num_tasks // num_nodes) + if max_mem_per_node_gb is not None: + memory_gb = min(memory_gb, max_mem_per_node_gb) + cluster_preset["mem"] = f"{memory_gb}G" # Build the slurm script script = SlurmScript( diff --git a/profiling/submit_poisson_weak_scaling.py b/profiling/submit_poisson_weak_scaling.py index f17ff7f4e..c05b46c8b 100644 --- a/profiling/submit_poisson_weak_scaling.py +++ b/profiling/submit_poisson_weak_scaling.py @@ -20,6 +20,13 @@ def main() -> None: default=64, help="Number of cells assigned to each rank in the weak-scaling direction.", ) + parser.add_argument( + "--ranks", + type=int, + nargs="+", + default=(1, 2, 4), + help="MPI rank counts to submit (default: 1 2 4).", + ) args = parser.parse_args() script_dir = Path(__file__).resolve().parent @@ -42,7 +49,7 @@ def main() -> None: param_flags = ["--cells-per-rank", str(args.cells_per_rank)] - for num_tasks in (1, 2, 4): #, 8, 16, 32, 64, 128, 256): + for num_tasks in args.ranks: profiling_case.launch(num_tasks, param_flags=param_flags) profiling_case.finalize_run() diff --git a/setup/modules.sh b/setup/modules.sh index 44691e231..bfdc8028f 100755 --- a/setup/modules.sh +++ b/setup/modules.sh @@ -91,6 +91,13 @@ case "$ACTION" in echo "Loading modules for $MACHINE, MODULES=$MODULES" module purge module load $MODULES + # Raven's Open MPI module intentionally does not populate LD_LIBRARY_PATH. + # Python extension modules such as mpi4py still need libmpi to be visible + # when it is loaded dynamically rather than linked through mpicc. + if [[ "$MACHINE" == "raven" && "$COMPILER_FAMILY" == "gcc" ]]; then + MPI_LIBDIR="$(mpicc --showme:libdirs)" + export LD_LIBRARY_PATH="$MPI_LIBDIR${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" + fi module list ;; display)