diff --git a/examples/Poisson/gaussian_cube/params_gaussian_cube.py b/examples/Poisson/gaussian_cube/params_gaussian_cube.py index d1fa7ce63..bd719774b 100644 --- a/examples/Poisson/gaussian_cube/params_gaussian_cube.py +++ b/examples/Poisson/gaussian_cube/params_gaussian_cube.py @@ -37,10 +37,41 @@ action="store_true", help="Small resolution (num_elements = [16, 16, 16]), for the tests (test_examples.py).", ) +parser.add_argument( + "--steps", + type=int, + default=1, + help="Number of time steps, i.e. Poisson solves with the same matrices (assembled once in the setup).", +) +parser.add_argument( + "--solver", + default="pcg", + choices=["pcg", "pcg_sr"], + help="CG variant: 'pcg_sr' does one global reduction per iteration instead of two (for many ranks/GPUs).", +) +parser.add_argument( + "--gpu-timing", + action="store_true", + help="Setup of the CPU vs GPU cases: one thread per rank, CUDA-event device times per region on the CuPy " + "backend instead of the line profiler, and a run label naming the backend and the rank count.", +) +parser.add_argument( + "--no-output", + action="store_true", + help="Write only the profiling data: no field or restart data (the pproc script then does nothing).", +) args, _ = parser.parse_known_args() if args.test_run: args.num_elements = [16, 16, 16] +if args.gpu_timing: + # One thread per MPI rank on both backends: OMP_NUM_THREADS=1, and the same for OpenBLAS/MKL, which would + # otherwise start one thread per core in every rank. Set before NumPy is imported (OpenBLAS reads it when loaded). + import os + + for _var in ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS"): + os.environ[_var] = "1" + precond = None if args.precond == "None" else args.precond # preconditioner -> (name suffix, linear solver) for the name and the description @@ -123,8 +154,8 @@ model = Poisson(base_units=base_units) # List all variables and decide whether to save their data -model.em_fields.phi.save_data = True -model.em_fields.source.save_data = True +model.em_fields.phi.save_data = not args.no_output +model.em_fields.source.save_data = not args.no_output # -------------------------- # Instance of the simulation @@ -135,10 +166,14 @@ out_folders=args.out_folder, sim_folder=args.sim_folder, restart=False, + save_step=args.steps, # save the first and the last snapshot only + save_restart=not args.no_output, ) -# Time stepping -time_opts = Time() +# Time stepping: with --steps > 1 every step solves the Poisson problem again (from x0 = 0, recycle=False) with the +# matrices assembled once in the setup, as in a time-dependent simulation; the solution is the same at every step. +dt = 0.01 +time_opts = Time() if args.steps == 1 else Time(dt=dt, Tend=round(args.steps * dt, 14)) # Geometry Lx = 2.0 @@ -155,10 +190,33 @@ # Derham options derham_opts = DerhamOptions(degree=(1, 2, 3), bcs=(("dirichlet", "dirichlet"), None, None)) -# Profilinig options -profiling_opts = ProfilingOptions( - use_line_profiler=True, -) +# Profiling options +if args.gpu_timing: + # CUDA-event device times per region on the CuPy backend (not on cunumpy's fake CuPy, whose CUDA events have no + # elapsed time). No line profiler: it would distort the CPU/GPU comparison. + import cunumpy as xp + from cunumpy.kernel_testing import fake_cupy_active + from maybempi import MPI, SerialMPI + + gpu_timing = xp.get_backend() == "cupy" and not fake_cupy_active() + + # scope-profiler run label, the name of the run in the reports and comparison plots, e.g. + # "CPU: 28 MPI ranks" and "GPU: 4 MPI ranks, 4 GPUs" + num_ranks = 1 if isinstance(MPI, SerialMPI) else MPI.COMM_WORLD.Get_size() + profiling_label = ( + f"{'GPU' if xp.get_backend() == 'cupy' else 'CPU'}: " + f"{num_ranks} MPI rank{'s' if num_ranks != 1 else ''}" + + (f", {num_ranks} GPU{'s' if num_ranks != 1 else ''}" if xp.get_backend() == "cupy" else "") + ) + profiling_opts = ProfilingOptions( + label=profiling_label, + use_gpu_timing=gpu_timing, + gpu_timing_backend="cupy", + ) +else: + profiling_opts = ProfilingOptions( + use_line_profiler=True, + ) # Simulation object sim = Simulation( @@ -184,7 +242,7 @@ solver_params = SolverParameters(tol=1e-8, maxiter=3000, info=True, recycle=False) model.propagators.poisson.options = model.propagators.poisson.Options( stab_eps=0.0, - solver="pcg", + solver=args.solver, precond=precond, solver_params=solver_params, ) @@ -236,4 +294,4 @@ def rhs_fun(e1, e2, e3): if __name__ == "__main__": # post-processing and plots: pproc_gaussian_cube.py - sim.run(profiling_activated=True, one_time_step=True) + sim.run(profiling_activated=True, one_time_step=args.steps == 1) diff --git a/examples/Poisson/gaussian_cube/pproc_gaussian_cube.py b/examples/Poisson/gaussian_cube/pproc_gaussian_cube.py index bfa46d7cf..b275e5b52 100644 --- a/examples/Poisson/gaussian_cube/pproc_gaussian_cube.py +++ b/examples/Poisson/gaussian_cube/pproc_gaussian_cube.py @@ -31,6 +31,10 @@ def main(path_out=DEFAULT_OUTPUT, save_figs=True, angle=np.pi): + if args.no_output: + print("The run wrote no field data (--no-output): nothing to post-process.") + return + out = Output(path_out).pproc(create_vtk=True, celldivide=CELLDIVIDE) # Raw FEEC fields are evaluated directly from the saved spline coefficients at the diff --git a/examples/Poisson/gaussian_cube/submit_gaussian_cube_cpu_vs_gpu.py b/examples/Poisson/gaussian_cube/submit_gaussian_cube_cpu_vs_gpu.py new file mode 100644 index 000000000..563ff565e --- /dev/null +++ b/examples/Poisson/gaussian_cube/submit_gaussian_cube_cpu_vs_gpu.py @@ -0,0 +1,108 @@ +"""Poisson on a 3D cube: 28 CPU MPI ranks vs 1 GPU, on the same kind of Pitagora booster node. + +This file defines the profiling case (the `ProfilingCase`) and submits its two runs of the +problem of `submit_gaussian_cube.py` (manufactured Gaussian solution on a Cuboid, +degree (1, 2, 3), unpreconditioned CG), at 128 x 128 x 128 cells and with 5 time steps, i.e. 5 +Poisson solves with the mass matrices assembled once in the setup, as in a time-dependent simulation: + +- CPU: 28 MPI ranks with one thread each (OMP_NUM_THREADS=1), NumPy backend; +- GPU: 1 MPI rank with one thread on one H100, CuPy backend. + +Every rank gets one CPU core (`cpus_per_task=1` of the preset), and the params file sets +OMP_NUM_THREADS=1 (and OPENBLAS/MKL_NUM_THREADS=1) before NumPy is imported (`--gpu-timing`). + +Both runs use the `pitagora_booster` preset, so they run on identical hardware; the CPU run does not +touch the GPUs. Without a batch system (or with STRUPHY_PROFILING_LOCAL=1) the runs execute directly +on this machine, where the 28-rank run is skipped unless there are that many cores. +`finalize_run` then packages and uploads each run as soon as its own job finishes. +Each generated script runs the simulation itself by invoking `params_gaussian_cube.py` +directly (its `__main__` block is the worker), and then post-processes it with `pproc_gaussian_cube.py` +(errors of the RHS and the solution, figures). + +Why this size: the GPU's advantage grows with the grid (the CG solve is bound by memory bandwidth, the +CPU's by its 28 cores); at 128^3 one solve takes 3.3 s on the GPU vs ~66 s on 28 login-node ranks, +and the setup (mass-matrix assembly in CUDA) 1.4 s vs ~15 s. + +The runs are profiled with scope-profiler, the CuPy run with CUDA-event timing per region; the +runs are labelled "CPU: 28 MPI ranks" and "GPU: 1 MPI rank, 1 GPU". +""" + +import argparse +import sys +from pathlib import Path + +# the profiling driver (profiling_job.py, clusters.py, ...) lives in profiling/ at the repo root, three levels up +sys.path.insert(0, str(Path(__file__).resolve().parents[3] / "profiling")) + +from profiling_job import ProfilingCase # noqa: E402 + + +def main() -> None: + + # Parse arguments, do not remove --upload + parser = argparse.ArgumentParser( + description=("Submit profiling jobs to a SLURM cluster and package the results for upload."), + ) + parser.add_argument( + "--upload", + action="store_true", + help="Upload the packaged profiling results to the profiling-data repo.", + ) + parser.add_argument( + "--language", + default="fortran", + choices=["fortran", "c"], + help="Pyccel language to compile the Struphy kernels with.", + ) + parser.add_argument( + "--compiler", + default="GNU", + help='Pyccel compiler family ("GNU", "intel", "PGI", "nvidia" or "LLVM").', + ) + args = parser.parse_args() + + # Paths relative to this script's location, so it can be run from anywhere. + script_dir = Path(__file__).resolve().parent + + profiling_case = ProfilingCase( + label="poisson_cube_cpu_node_vs_gpu", + name="Poisson on 3D cube: 28 CPU ranks vs 1 GPU", + description="The Poisson model with manufactured solution on a 3D cube (Cuboid, 128^3 cells, degree (1, 2, 3)), 5 solves with unpreconditioned CG, on 28 MPI ranks x 1 thread of a booster node (NumPy) vs 1 MPI rank x 1 thread on 1 H100 GPU (CuPy).", + physics_problem="Occurs in many plasma applications.", + struphy_model_used="Poisson", + params_source=script_dir / "params_gaussian_cube.py", + pproc_source=script_dir / "pproc_gaussian_cube.py", + language=args.language, + compiler=args.compiler, + upload=args.upload, + ) + + # 128^3 cells, 5 solves (see the module docstring); the CPU run takes a few minutes + param_flags = ["--num-elements", "128", "128", "128", "--steps", "5", "--gpu-timing"] + slurm_options = {"time": "00:30:00"} + + # CPU: 28 MPI ranks x 1 thread on a booster node, NumPy backend + profiling_case.launch( + 28, + num_nodes=1, + param_flags=param_flags, + backend="numpy", + slurm_preset="pitagora_booster", + slurm_options=slurm_options, + ) + # GPU: 1 MPI rank x 1 thread on one H100 (the rank binds to GPU 0 of the node), CuPy backend + profiling_case.launch( + 1, + num_nodes=1, + param_flags=param_flags, + backend="cupy", + slurm_preset="pitagora_booster", + slurm_options=slurm_options, + ) + + # Package and push each run as its own job finishes. + profiling_case.finalize_run() + + +if __name__ == "__main__": + main() diff --git a/examples/Poisson/gaussian_cube/submit_gaussian_cube_gpu_scaling.py b/examples/Poisson/gaussian_cube/submit_gaussian_cube_gpu_scaling.py new file mode 100644 index 000000000..356f435b1 --- /dev/null +++ b/examples/Poisson/gaussian_cube/submit_gaussian_cube_gpu_scaling.py @@ -0,0 +1,117 @@ +"""Poisson on a 3D cube: strong scaling on 1, 2 and 4 GPUs of a Pitagora booster node. + +This file defines the profiling case (the `ProfilingCase`) and submits its three runs of the same +problem as `submit_gaussian_cube_cpu_vs_gpu.py` (manufactured Gaussian solution on a Cuboid, +degree (1, 2, 3), unpreconditioned CG, 5 solves), by default on 192 x 192 x 192 cells (`--cells`), on the +CuPy backend with 1, 2 and 4 MPI ranks (`--gpus`), one GPU per rank (feectools binds each rank to the GPU of its +node-local rank), one thread per rank. `--solver pcg_sr` selects the CG variant with one global reduction +per iteration. The runs write only their profiling data (`--no-output` of the params file), so there is no +post-processing. Each run uses the `pitagora_booster` preset (one node, its 4 GPUs). +Without a batch system (or with STRUPHY_PROFILING_LOCAL=1) the runs execute directly on this machine. +`finalize_run` then packages and uploads each run as soon as its own job finishes. +Each generated script runs the simulation itself by invoking `params_gaussian_cube.py` +directly (its `__main__` block is the worker). + +The runs are profiled with scope-profiler (CUDA-event timing per region) and labelled +"GPU: N MPI ranks, N GPUs". The case label contains the size (and the solver if not pcg), e.g. +`poisson_cube_gpu_scaling_192`. + +Recommended runs: the default (192^3 on 1, 2 and 4 GPUs) for the strong scaling from one GPU, and +`--cells 256 --gpus 2 4` for a larger problem (256^3 does not fit on one H100: the setup assembles M0 and M1 as +3D stencil matrices and evaluates the metric at all quadrature points at once). The fixed cost per CG +iteration (ghost exchanges, reductions) weighs less on the larger problem: 4 vs 2 GPUs gave 1.30x at 192^3 and +1.61x at 256^3 (Oct 2026). +""" + +import argparse +import sys +from pathlib import Path + +# the profiling driver (profiling_job.py, clusters.py, ...) lives in profiling/ at the repo root, three levels up +sys.path.insert(0, str(Path(__file__).resolve().parents[3] / "profiling")) + +from profiling_job import ProfilingCase # noqa: E402 + + +def main() -> None: + + # Parse arguments, do not remove --upload + parser = argparse.ArgumentParser( + description=("Submit profiling jobs to a SLURM cluster and package the results for upload."), + ) + parser.add_argument( + "--upload", + action="store_true", + help="Upload the packaged profiling results to the profiling-data repo.", + ) + parser.add_argument( + "--language", + default="fortran", + choices=["fortran", "c"], + help="Pyccel language to compile the Struphy kernels with.", + ) + parser.add_argument( + "--compiler", + default="GNU", + help='Pyccel compiler family ("GNU", "intel", "PGI", "nvidia" or "LLVM").', + ) + parser.add_argument( + "--cells", + type=int, + default=192, + help="Number of cells in each direction. 192^3 fits on one H100; 256^3 needs at least 2 (--gpus 2 4).", + ) + parser.add_argument( + "--gpus", + type=int, + nargs="+", + default=[1, 2, 4], + choices=[1, 2, 3, 4], + help="GPU counts (MPI ranks, one GPU each) to run, on one booster node.", + ) + parser.add_argument( + "--solver", + default="pcg", + choices=["pcg", "pcg_sr"], + help="CG variant: 'pcg_sr' does one global reduction per iteration instead of two.", + ) + args = parser.parse_args() + + # Paths relative to this script's location, so it can be run from anywhere. + script_dir = Path(__file__).resolve().parent + gpus = sorted(set(args.gpus)) + gpus_text = ", ".join(str(n) for n in gpus[:-1]) + (" and " if len(gpus) > 1 else "") + str(gpus[-1]) + + profiling_case = ProfilingCase( + # one label per size and solver: runs of different sizes are separate cases in the profiling-data repo + label=f"poisson_cube_gpu_scaling_{args.cells}" + ("" if args.solver == "pcg" else f"_{args.solver}"), + name=f"Poisson on 3D cube, {args.cells}^3 cells: strong scaling on {gpus_text} GPUs", + description=f"The Poisson model with manufactured solution on a 3D cube (Cuboid, {args.cells}^3 cells, degree (1, 2, 3)), 5 solves with unpreconditioned CG ('{args.solver}'), on {gpus_text} H100 GPUs of a booster node (CuPy, one MPI rank per GPU).", + physics_problem="Occurs in many plasma applications.", + struphy_model_used="Poisson", + params_source=script_dir / "params_gaussian_cube.py", + language=args.language, + compiler=args.compiler, + upload=args.upload, + ) + + # 5 solves, as in submit_gaussian_cube_cpu_vs_gpu.py; only the profiling data is written + cells = [str(args.cells)] * 3 + param_flags = ["--num-elements", *cells, "--steps", "5", "--solver", args.solver, "--gpu-timing", "--no-output"] + + # one MPI rank per GPU, on one booster node + for num_tasks in gpus: + profiling_case.launch( + num_tasks, + num_nodes=1, + param_flags=param_flags, + backend="cupy", + slurm_preset="pitagora_booster", + ) + + # Package and push each run as its own job finishes. + profiling_case.finalize_run() + + +if __name__ == "__main__": + main() diff --git a/examples/Poisson/sweep_poisson.py b/examples/Poisson/sweep_poisson.py new file mode 100644 index 000000000..412d8e09d --- /dev/null +++ b/examples/Poisson/sweep_poisson.py @@ -0,0 +1,72 @@ +"""Run all Poisson profiling cases of examples/Poisson as one sweep (see profiling/sweep.py). + +- CPU strong scaling (1 to 256 MPI ranks), every variant of every case: + gaussian_cube, gaussian_hollow_cylinder, gaussian_hollow_torus (preconditioners), + gaussian_on_axis_cylinder, gaussian_on_axis_torus (tensor-product vs polar splines, preconditioners). + The variants are read from the VARIANTS dict of each submit script, so new variants are picked up. +- GPU cases of gaussian_cube: 28 CPU ranks vs 1 GPU (128^3), and GPU strong scaling at 192^3 (1, 2, 4 GPUs) and + 256^3 (2, 4 GPUs; 256^3 does not fit on one GPU). + +All jobs are submitted first (SLURM runs them as the partitions allow); the runs are then packaged, and with +--upload pushed, one at a time as their jobs finish. Keep this process running until it is done (e.g. in tmux). + +Usage (from anywhere, with the struphy venv activated): + + python examples/Poisson/sweep_poisson.py [--upload] [--cpu-only | --gpu-only] [--dry-run] +""" + +import argparse +import sys +from pathlib import Path + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE.parents[1] / "profiling")) + +from sweep import run_sweep, variant_entries # noqa: E402 + +CPU_CASES = ( + "gaussian_cube", + "gaussian_hollow_cylinder", + "gaussian_hollow_torus", + "gaussian_on_axis_cylinder", + "gaussian_on_axis_torus", +) + + +def entries(cpu: bool = True, gpu: bool = True) -> list: + """The submit commands of the sweep, as (script, arguments).""" + out = [] + if cpu: + for case in CPU_CASES: + out += variant_entries(HERE / case / f"submit_{case}.py") + if gpu: + cube = HERE / "gaussian_cube" + out += [ + (cube / "submit_gaussian_cube_cpu_vs_gpu.py", []), + (cube / "submit_gaussian_cube_gpu_scaling.py", ["--cells", "192", "--gpus", "1", "2", "4"]), + (cube / "submit_gaussian_cube_gpu_scaling.py", ["--cells", "256", "--gpus", "2", "4"]), + ] + return out + + +def main() -> None: + parser = argparse.ArgumentParser(description="Run all Poisson profiling cases as one sweep.") + parser.add_argument("--upload", action="store_true", help="Push the packaged runs to the profiling-data repo.") + which = parser.add_mutually_exclusive_group() + which.add_argument("--cpu-only", action="store_true", help="Only the CPU strong scaling cases.") + which.add_argument("--gpu-only", action="store_true", help="Only the GPU cases of gaussian_cube.") + parser.add_argument("--dry-run", action="store_true", help="Only print the commands.") + parser.add_argument("--poll-interval", type=float, default=30.0, help="Seconds between polls of the SLURM queue.") + args = parser.parse_args() + + n_failed = run_sweep( + entries(cpu=not args.gpu_only, gpu=not args.cpu_only), + upload=args.upload, + dry_run=args.dry_run, + poll_interval=args.poll_interval, + ) + sys.exit(1 if n_failed else 0) + + +if __name__ == "__main__": + main() diff --git a/examples/Vlasov/homogeneous_slab/params_homogeneous_slab.py b/examples/Vlasov/homogeneous_slab/params_homogeneous_slab.py new file mode 100644 index 000000000..1234bca97 --- /dev/null +++ b/examples/Vlasov/homogeneous_slab/params_homogeneous_slab.py @@ -0,0 +1,231 @@ +"""Particle pushing in a homogeneous slab (Vlasov model, no field solve): a CPU vs GPU benchmark. The pproc writes +the scope-profiler report and the cost per step and marker.""" + +# ----------------------------- +# 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 = "Vlasov in a homogeneous slab: 28 CPU ranks vs 1 GPU" +description = """ +The Vlasov model (markers pushed in a static, homogeneous background magnetic field, no field solve): a pure +particle-pushing benchmark. Unit cube, periodic boundaries, HomogenSlab equilibrium (constant B), a Maxwellian +with a cosine density perturbation plus a second, uniform Maxwellian, 16 million markers, 20 steps. +Each step runs the PushVxB (rotation of v about B) and PushEta (position update, explicit Runge-Kutta) +propagators, followed by the MPI exchange of markers that left their process domain. + +Runs on the NumPy (CPU) or CuPy (GPU) backend, chosen by the environment variable CUNUMPY_BACKEND (the profiling +driver exports it in the job script). All timing comes from scope-profiler; on the CuPy backend it records +CUDA-event device time per region. The post-processing (pproc_homogeneous_slab.py) writes the scope-profiler HTML +report and a summary of the cost per step and marker into the run's results folder. +""" + +# One thread per MPI rank on both backends (CPU: 28 ranks x 1 thread, GPU: 1 rank x 1 thread): OMP_NUM_THREADS=1, and +# the same for OpenBLAS/MKL, which would otherwise start one thread per core in every rank. Set unconditionally, before +# NumPy is imported (OpenBLAS reads it when it is loaded). +import os + +for _var in ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS"): + os.environ[_var] = "1" + +import argparse +import logging +from pathlib import Path + +from struphy import set_logging_level + +set_logging_level(logging.INFO) + +import cunumpy as xp + +# ------------------ +# Import Struphy API +# ------------------ +from struphy import ( + BaseUnits, + BoundaryParameters, + DerhamOptions, + EnvironmentOptions, + LoadingParameters, + ProfilingOptions, + SavingParameters, + Simulation, + SortingParameters, + Time, + WeightsParameters, + domains, + equils, + grids, + maxwellians, + perturbations, +) + +# -------------------------- +# Command line parameters +# -------------------------- +# `--out-folder` and `--sim-folder` set the output folder; the profiling driver passes its own +# (see `ProfilingCase.build_commands`). pproc_homogeneous_slab.py takes the same flags. +# Unknown flags are ignored so the driver can forward other parameters as well. +# no -h here when imported by the pproc script, which has its own help +parser = argparse.ArgumentParser(add_help=__name__ == "__main__") +parser.add_argument( + "--out-folder", + default=str(Path(__file__).resolve().parent), + help="Folder in which the output folder is created (default: beside this file).", +) +parser.add_argument("--sim-folder", default="sim_1", help="Name of the output folder.") +parser.add_argument("--n-markers", type=int, default=16_000_000, help="Total number of markers.") +parser.add_argument("--steps", type=int, default=20, help="Number of time steps.") +parser.add_argument( + "--test-run", + action="store_true", + help="10,000 markers and 2 steps, for the tests (test_examples.py).", +) +args, _ = parser.parse_known_args() +if args.test_run: + args.n_markers = 10_000 + args.steps = 2 + +# --------------------- +# Instance of the model +# --------------------- +from struphy.models import Vlasov + +# Units +base_units = BaseUnits() + +# Model instance +model = Vlasov(base_units=base_units) + +# List all variables and decide whether to save their data +model.kinetic_ions.var.save_data = False # profile the pushing, not the marker output + +# -------------------------- +# Instance of the simulation +# -------------------------- + +# Environment options +env = EnvironmentOptions( + out_folders=args.out_folder, + sim_folder=args.sim_folder, + restart=False, + save_restart=False, # restart data would dominate the run time for many markers +) + +# Time stepping +dt = 0.01 +steps = args.steps +time_opts = Time(dt=dt, Tend=round(steps * dt, 14)) + +# Geometry +domain = domains.Cuboid() + +# Fluid equilibrium: constant background magnetic field +equil = equils.HomogenSlab() + +# Grid +# The grid only carries the projected background field and defines the MPI domains of the markers. 28 x 16 x 16 +# cells, decomposed in all directions: 28 ranks get a 7 x 2 x 2 process grid, i.e. a 4 x 8 x 8 block of cells each. +grid = grids.TensorProductGrid(num_elements=(28, 16, 16), mpi_dims_mask=(True, True, True)) + +# Derham options +derham_opts = DerhamOptions() + +# Profiling options: CUDA-event device times per region on the CuPy backend. +# cunumpy's fake CuPy (CUNUMPY_FAKE_CUPY=1, no GPU) has CUDA events without elapsed time, so GPU timing needs a GPU. +from cunumpy.kernel_testing import fake_cupy_active + +gpu_timing = xp.get_backend() == "cupy" and not fake_cupy_active() + +# scope-profiler run label, the name of the run in the reports and comparison plots, e.g. +# "CPU (numpy): 28 MPI ranks x 1 thread" and "GPU (cupy): 1 MPI rank x 1 thread, 1 GPU" +from maybempi import MPI, SerialMPI + +num_ranks = 1 if isinstance(MPI, SerialMPI) else MPI.COMM_WORLD.Get_size() +num_threads = int(os.environ["OMP_NUM_THREADS"]) +profiling_label = ( + f"{'GPU' if xp.get_backend() == 'cupy' else 'CPU'} ({xp.get_backend()}): " + f"{num_ranks} MPI rank{'s' if num_ranks != 1 else ''} x {num_threads} thread{'s' if num_threads != 1 else ''}" + + (f", {num_ranks} GPU{'s' if num_ranks != 1 else ''}" if xp.get_backend() == "cupy" else "") +) + +profiling_opts = ProfilingOptions( + label=profiling_label, + use_gpu_timing=gpu_timing, + gpu_timing_backend="cupy", +) + +# 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, + profiling_opts=profiling_opts, +) + +# ------------------- +# Particle parameters +# ------------------- + +# 16 million markers in total: 4 GB of marker data on the GPU, ~570k markers per rank on 28 CPU ranks. +# pseudo_random loading is fast on both backends (sobol_standard is a Python loop, ~1 s per 100k markers per rank); +# NumPy and CuPy draw different markers from the same seed, so the two runs agree statistically, not value by value. +n_markers = args.n_markers +loading_params = LoadingParameters(Np=n_markers, seed=1234) +weights_params = WeightsParameters() +boundary_params = BoundaryParameters() +# sorting boxes are off: SortingBoxes does not work on the CuPy backend yet (struphy-hub/struphy#726) +sorting_params = SortingParameters() +saving_params = SavingParameters(n_markers=0) # profile the pushing, not the marker output + +model.kinetic_ions.set_markers( + loading_params=loading_params, + weights_params=weights_params, + boundary_params=boundary_params, + sorting_params=sorting_params, + saving_params=saving_params, + # room for markers arriving from other ranks: 25% of the markers (default 100%). Every marker kernel and column + # operation runs over all rows of the buffer, holes included, so a smaller buffer makes each step cheaper on both + # backends; with 28 ranks a rank's domain is 1/7 of the box in x and markers move ~0.01 per step. + bufsize=0.25, +) + +# ------------------ +# Propagator options +# ------------------ + +model.propagators.push_vxb.options = model.propagators.push_vxb.Options() +model.propagators.push_eta.options = model.propagators.push_eta.Options() + +# ------------------ +# Initial conditions +# ------------------ +# For kinetic species the background is mandatory; the perturbation is added to the density. + +maxwellian_1 = maxwellians.Maxwellian3D(n=(1.0, None)) +maxwellian_2 = maxwellians.Maxwellian3D(n=(0.1, None)) +model.kinetic_ions.var.add_background(maxwellian_1 + maxwellian_2) + +perturbation = perturbations.TorusModesCos() +model.kinetic_ions.var.add_initial_condition(maxwellians.Maxwellian3D(n=(1.0, perturbation)) + maxwellian_2) + + +if __name__ == "__main__": + import contextlib + + # Without a GPU, CUNUMPY_FAKE_CUPY=1 CUNUMPY_BACKEND=cupy runs the CuPy code path on cunumpy's fake CuPy with every + # CUDA launch emulated on the CPU: a check that the model runs on the device backend, not a timing. On a GPU + # (or on NumPy) the context does nothing. + from struphy.models.tests.gpu_e2e import device_context + + # post-processing (profiling report and cost summary): pproc_homogeneous_slab.py + with device_context() if fake_cupy_active() else contextlib.nullcontext(): + sim.run(profiling_activated=True) diff --git a/examples/Vlasov/homogeneous_slab/pproc_homogeneous_slab.py b/examples/Vlasov/homogeneous_slab/pproc_homogeneous_slab.py new file mode 100644 index 000000000..4fbfc9801 --- /dev/null +++ b/examples/Vlasov/homogeneous_slab/pproc_homogeneous_slab.py @@ -0,0 +1,94 @@ +"""Post-process the Vlasov homogeneous_slab benchmark. + +Run as ``python pproc_homogeneous_slab.py [sim_folder] [flags of params_homogeneous_slab.py]``; the flags must be +those of the run. Writes the scope-profiler HTML report and a summary of the cost per step and marker into +``/results``, where the profiling packaging picks them up and uploads them with the run. +``submit_homogeneous_slab_cpu_vs_gpu.py`` and ``submit_homogeneous_slab_gpu_scaling.py`` call this script after each +simulation. +""" + +import argparse +import json +import statistics +import sys +from pathlib import Path + +import cunumpy as xp +import numpy as np + +from struphy import Output + +SCRIPT_DIR = Path(__file__).resolve().parent +DEFAULT_OUTPUT = SCRIPT_DIR / "sim_1" + +# The setup of the run (parsed from the same command line). +sys.path.insert(0, str(SCRIPT_DIR)) +import params_homogeneous_slab as params # noqa: E402 + + +def main(path_out=DEFAULT_OUTPUT): + from scope_profiler import create_html_report + + if not (Path(path_out) / "profiling_data.h5").exists(): + print(f"No profiling data in {path_out} (run without profiling): nothing to write.") + return + + out = Output(path_out) + results = out.profile.results # the scope-profiler results of the run, all ranks merged + results_dir = out.path_out / "results" + results_dir.mkdir(exist_ok=True) + + # the full scope-profiler report: call tree, charts, per-rank statistics + create_html_report(results, results_dir / "report.html") + + # cost per step and per marker, from the `model.integrate` region (the maximum over the ranks per step); + # the first step includes one-time work (CUDA module loads, allocations of work arrays). + # On the CuPy backend the CPU timestamps only measure kernel launches (no device sync at region ends), so + # the CUDA-event device times are used whenever they were recorded. + integrate = results["model.integrate"] + per_rank = [ + integrate[rank].gpu_durations if integrate[rank].has_gpu_timing else integrate[rank].durations + for rank in integrate.ranks + ] + n_calls = min(len(durations) for durations in per_rank) + step_times = np.max([durations[:n_calls] for durations in per_rank], axis=0) + step_median = float(statistics.median(step_times[1:] if len(step_times) > 1 else step_times)) + summary = { + "backend": xp.get_backend(), + "gpu_timing": params.gpu_timing, + "step_times_from": "gpu" if integrate[integrate.ranks[0]].has_gpu_timing else "cpu", + "mpi_ranks": results.num_ranks, + "num_elements": list(params.sim.grid.num_elements), + "n_markers": params.n_markers, + "steps": params.steps, + "dt": params.dt, + "step_times": step_times.tolist(), + "step_median": step_median, + "ns_per_marker_step": 1e9 * step_median / params.n_markers, + } + (results_dir / "summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8") + + print( + f"\nstep: median {1e3 * step_median:.3f} ms (first step {1e3 * step_times[0]:.3f} ms), " + f"{summary['ns_per_marker_step']:.1f} ns per marker and step ({params.n_markers:,} markers)" + ) + # the propagators and kernels of the run; on the CuPy backend with GPU columns + results.print_summary( + include=["^prop: ", "^kernel: ", "^pusher: ", "mpi_sort"], + sort="total", + title="propagators and kernels", + ) + + +if __name__ == "__main__": + # the flags of the run are parsed by params_homogeneous_slab.py, the rest here + _, other_flags = params.parser.parse_known_args() + parser = argparse.ArgumentParser(description="Write the profiling results of a saved simulation run.") + parser.add_argument( + "path_out", + nargs="?", + default=DEFAULT_OUTPUT, + help="Simulation output folder (default: sim_1 beside this script)", + ) + args = parser.parse_args(other_flags) + main(args.path_out) diff --git a/examples/Vlasov/homogeneous_slab/submit_homogeneous_slab_cpu_vs_gpu.py b/examples/Vlasov/homogeneous_slab/submit_homogeneous_slab_cpu_vs_gpu.py new file mode 100644 index 000000000..16be0d66c --- /dev/null +++ b/examples/Vlasov/homogeneous_slab/submit_homogeneous_slab_cpu_vs_gpu.py @@ -0,0 +1,86 @@ +"""Vlasov in a homogeneous slab: 28 CPU MPI ranks vs 1 GPU, on the same kind of Pitagora booster node. + +This file defines the profiling case (the `ProfilingCase`) and submits its two runs of the same +problem, a pure particle-pushing benchmark with the Vlasov model (16 million markers, PushVxB and +PushEta in a constant background magnetic field, 20 steps, no field solve): + +- CPU: 28 MPI ranks with one thread each (OMP_NUM_THREADS=1), NumPy backend; +- GPU: 1 MPI rank with one thread on one H100, CuPy backend. + +Every rank gets one CPU core (`cpus_per_task=1` of the preset), and the params file sets +OMP_NUM_THREADS=1 (and OPENBLAS/MKL_NUM_THREADS=1) before NumPy is imported. + +Both runs use the `pitagora_booster` preset, so they run on identical hardware; the CPU run does not +touch the GPUs. Without a batch system (or with STRUPHY_PROFILING_LOCAL=1) the runs execute directly +on this machine, where the 28-rank run is skipped unless there are that many cores. +`finalize_run` then packages and uploads each run as soon as its own job finishes. +Each generated script runs the simulation itself by invoking `params_homogeneous_slab.py` +directly (its `__main__` block is the worker), and then writes the scope-profiler report and the cost +summary with `pproc_homogeneous_slab.py`. + +The runs are profiled with scope-profiler, the CuPy run with CUDA-event timing per region; the +runs are labelled "CPU (numpy): 28 MPI ranks x 1 thread" and "GPU (cupy): 1 MPI rank x 1 thread, 1 GPU". +The resolution (28 x 16 x 16 cells, 16M markers, 20 steps) is set in `params_homogeneous_slab.py`. +""" + +import argparse +import sys +from pathlib import Path + +# the profiling driver (profiling_job.py, clusters.py, ...) lives in profiling/ at the repo root, three levels up +sys.path.insert(0, str(Path(__file__).resolve().parents[3] / "profiling")) + +from profiling_job import ProfilingCase # noqa: E402 + + +def main() -> None: + + # Parse arguments, do not remove --upload + parser = argparse.ArgumentParser( + description=("Submit profiling jobs to a SLURM cluster and package the results for upload."), + ) + parser.add_argument( + "--upload", + action="store_true", + help="Upload the packaged profiling results to the profiling-data repo.", + ) + parser.add_argument( + "--language", + default="fortran", + choices=["fortran", "c"], + help="Pyccel language to compile the Struphy kernels with.", + ) + parser.add_argument( + "--compiler", + default="GNU", + help='Pyccel compiler family ("GNU", "intel", "PGI", "nvidia" or "LLVM").', + ) + args = parser.parse_args() + + # Paths relative to this script's location, so it can be run from anywhere. + script_dir = Path(__file__).resolve().parent + + profiling_case = ProfilingCase( + label="vlasov_homogeneous_slab_cpu_node_vs_gpu", + name="Vlasov in a homogeneous slab: 28 CPU ranks vs 1 GPU", + description="The Vlasov model (16M markers pushed in a constant background magnetic field, PushVxB and PushEta, 20 steps, no field solve) on 28 MPI ranks x 1 thread of a booster node (NumPy) vs 1 MPI rank x 1 thread on 1 H100 GPU (CuPy).", + physics_problem="Charged particles gyrating and streaming in a constant magnetic field (pure particle pushing).", + struphy_model_used="Vlasov", + params_source=script_dir / "params_homogeneous_slab.py", + pproc_source=script_dir / "pproc_homogeneous_slab.py", + language=args.language, + compiler=args.compiler, + upload=args.upload, + ) + + # CPU: 28 MPI ranks x 1 thread on a booster node, NumPy backend + profiling_case.launch(28, num_nodes=1, backend="numpy", slurm_preset="pitagora_booster") + # GPU: 1 MPI rank x 1 thread on one H100 (the rank binds to GPU 0 of the node), CuPy backend + profiling_case.launch(1, num_nodes=1, backend="cupy", slurm_preset="pitagora_booster") + + # Package and push each run as its own job finishes. + profiling_case.finalize_run() + + +if __name__ == "__main__": + main() diff --git a/examples/Vlasov/homogeneous_slab/submit_homogeneous_slab_gpu_scaling.py b/examples/Vlasov/homogeneous_slab/submit_homogeneous_slab_gpu_scaling.py new file mode 100644 index 000000000..8baf55f4d --- /dev/null +++ b/examples/Vlasov/homogeneous_slab/submit_homogeneous_slab_gpu_scaling.py @@ -0,0 +1,99 @@ +"""Vlasov in a homogeneous slab: strong scaling on 1, 2 and 4 GPUs of a Pitagora booster node. + +This file defines the profiling case (the `ProfilingCase`) and submits its three runs of the same +problem as `submit_homogeneous_slab_cpu_vs_gpu.py` (16 million markers by default, `--n-markers`, PushVxB and +PushEta in a constant background magnetic field, 20 steps, no field solve), on the CuPy backend with +1, 2 and 4 MPI ranks (`--gpus`), one GPU per rank (feectools binds each rank to the GPU of its node-local rank), +one thread per rank. Each run uses the `pitagora_booster` preset (one node, its 4 GPUs). +Without a batch system (or with STRUPHY_PROFILING_LOCAL=1) the runs execute directly on this machine. +`finalize_run` then packages and uploads each run as soon as its own job finishes. +Each generated script runs the simulation itself by invoking `params_homogeneous_slab.py` +directly (its `__main__` block is the worker), and then writes the scope-profiler report and the cost +summary with `pproc_homogeneous_slab.py`. + +The runs are profiled with scope-profiler (CUDA-event timing per region) and labelled +"GPU (cupy): N MPI ranks x 1 thread, N GPUs". With pseudo-random loading every rank draws the same +global sequence and keeps the markers of its domain, so all three runs push the same markers. +""" + +import argparse +import sys +from pathlib import Path + +# the profiling driver (profiling_job.py, clusters.py, ...) lives in profiling/ at the repo root, three levels up +sys.path.insert(0, str(Path(__file__).resolve().parents[3] / "profiling")) + +from profiling_job import ProfilingCase # noqa: E402 + + +def main() -> None: + + # Parse arguments, do not remove --upload + parser = argparse.ArgumentParser( + description=("Submit profiling jobs to a SLURM cluster and package the results for upload."), + ) + parser.add_argument( + "--upload", + action="store_true", + help="Upload the packaged profiling results to the profiling-data repo.", + ) + parser.add_argument( + "--language", + default="fortran", + choices=["fortran", "c"], + help="Pyccel language to compile the Struphy kernels with.", + ) + parser.add_argument( + "--compiler", + default="GNU", + help='Pyccel compiler family ("GNU", "intel", "PGI", "nvidia" or "LLVM").', + ) + parser.add_argument( + "--n-markers", + type=int, + default=16_000_000, + help="Total number of markers (the same on every GPU count).", + ) + parser.add_argument( + "--gpus", + type=int, + nargs="+", + default=[1, 2, 4], + choices=[1, 2, 3, 4], + help="GPU counts (MPI ranks, one GPU each) to run, on one booster node.", + ) + args = parser.parse_args() + gpus = sorted(set(args.gpus)) + gpus_text = ", ".join(str(n) for n in gpus[:-1]) + (" and " if len(gpus) > 1 else "") + str(gpus[-1]) + millions = f"{args.n_markers / 1e6:g}M" + + # Paths relative to this script's location, so it can be run from anywhere. + script_dir = Path(__file__).resolve().parent + + profiling_case = ProfilingCase( + # the default size keeps the label of earlier runs; other sizes are separate cases + label="vlasov_homogeneous_slab_gpu_scaling" + ("" if args.n_markers == 16_000_000 else f"_{millions}"), + name=f"Vlasov in a homogeneous slab, {millions} markers: strong scaling on {gpus_text} GPUs", + description=f"The Vlasov model ({millions} markers pushed in a constant background magnetic field, PushVxB and PushEta, 20 steps, no field solve) on {gpus_text} H100 GPUs of a booster node (CuPy, one MPI rank per GPU).", + physics_problem="Charged particles gyrating and streaming in a constant magnetic field (pure particle pushing).", + struphy_model_used="Vlasov", + params_source=script_dir / "params_homogeneous_slab.py", + pproc_source=script_dir / "pproc_homogeneous_slab.py", + language=args.language, + compiler=args.compiler, + upload=args.upload, + ) + + # 1, 2 and 4 MPI ranks, one GPU each, on one booster node + param_flags = ["--n-markers", str(args.n_markers)] + for num_tasks in gpus: + profiling_case.launch( + num_tasks, num_nodes=1, param_flags=param_flags, backend="cupy", slurm_preset="pitagora_booster" + ) + + # Package and push each run as its own job finishes. + profiling_case.finalize_run() + + +if __name__ == "__main__": + main() diff --git a/profiling/profiling_job.py b/profiling/profiling_job.py index 3b1192ba0..2a25e3d03 100644 --- a/profiling/profiling_job.py +++ b/profiling/profiling_job.py @@ -9,9 +9,10 @@ a single object to thread through the rest of the run. - For each rank count, the caller calls `ProfilingCase.launch(ntasks, param_flags=None, backend="numpy")`, which builds the per-rank shell commands and either submits a `SlurmScript` (under SLURM) or - writes and runs a plain bash script to completion (otherwise) — the caller doesn't - need to know which. Pass `param_flags` to override the default commands (e.g. to add a - case-specific flag such as an arbitrary `ppc`) before they're wrapped in a script. + writes and runs a plain bash script to completion (otherwise, or with + STRUPHY_PROFILING_LOCAL=1) — the caller doesn't need to know which. Pass `param_flags` + to override the default commands (e.g. to add a case-specific flag such as an arbitrary + `ppc`) before they're wrapped in a script. Under SLURM, the preset of the detected cluster is taken from `clusters.SLURM_PRESETS` on each call, unless `slurm_preset` names another one (e.g. a GPU partition); `slurm_options` overrides single pragmas of it. @@ -80,9 +81,13 @@ def has_module_system() -> bool: ) -def detect_launcher() -> str: - """The MPI launcher to use: `srun` under SLURM, otherwise `mpirun`/`mpiexec`.""" - if shutil.which("srun"): +def detect_launcher(use_slurm: bool = True) -> str: + """The MPI launcher to use: `srun` under SLURM, otherwise `mpirun`/`mpiexec`. + + With `use_slurm=False` (local runs) `srun` is never used, even where it exists: on a login + node it would request a new allocation instead of running on this machine. + """ + if use_slurm and shutil.which("srun"): return "srun" for launcher in ("mpirun", "mpiexec"): if shutil.which(launcher): @@ -379,13 +384,17 @@ def setup_run(self) -> None: # pushing a run is just a commit — there is no intermediate export folder to # copy out of. Without `--upload` the clone is pointless, so the packaged # folders are simply written into the same, plain directory. + # Prepared once per process: the cases of a sweep (see `begin_sweep`) share the folder, and + # preparing it again (deleting or re-cloning it) would drop what the others packaged. self.output_root = profiling_results_base / "profiling-data" - if self.upload: - _clone_profiling_data(self.output_root) - else: - if self.output_root.exists(): - shutil.rmtree(self.output_root) - self.output_root.mkdir(parents=True, exist_ok=True) + if self.output_root not in _PREPARED_OUTPUT_ROOTS: + if self.upload: + _clone_profiling_data(self.output_root) + else: + if self.output_root.exists(): + shutil.rmtree(self.output_root) + self.output_root.mkdir(parents=True, exist_ok=True) + _PREPARED_OUTPUT_ROOTS.add(self.output_root) # Determine the current git commit hash for the Struphy repo self.run_commit = _git_commit(repo_root) @@ -402,9 +411,11 @@ def setup_run(self) -> None: self.case_output_root = self.run_results_root / self.label self.case_output_root.mkdir(parents=True, exist_ok=True) - self.use_slurm = shutil.which("sbatch") is not None + # STRUPHY_PROFILING_LOCAL=1 runs on this machine even where SLURM is available (e.g. on a login node) + force_local = os.environ.get("STRUPHY_PROFILING_LOCAL", "0").lower() in ("1", "true", "yes") + self.use_slurm = shutil.which("sbatch") is not None and not force_local self.use_modules = self.use_slurm and has_module_system() - self.launcher = detect_launcher() + self.launcher = detect_launcher(self.use_slurm) def case_info_dict(self) -> dict: """This case's metadata plus every submitted/launched job, as written to `profiling_case_info.json`. @@ -657,45 +668,64 @@ def finalize_run(self, poll_interval: float = 10.0) -> None: Whether the pushes happen at all is the case's `upload` flag, set at construction. + Inside a sweep (see `sweep.py` and `begin_sweep`), this only registers the case: + `finalize_all` then waits on the runs of all cases of the sweep together. + Args: poll_interval: Seconds between polls for finished runs. """ - print( - f"Writing metadata for '{self.label}' to {self.case_output_root / 'profiling_case_info.json'}", - ) - case_info = self.case_info_dict() - (self.case_output_root / "profiling_case_info.json").write_text( - json.dumps(case_info, indent=2), - encoding="utf-8", - ) + if _SWEEP_CASES is not None: + _SWEEP_CASES.append(self) + print(f"Submitted {len(self.job_infos)} run(s) of '{self.label}'; packaging deferred to the sweep.") + return - # Package and push the case-level metadata before waiting on any run, so the - # packaged folder is already in place when the per-run pushes start updating it. - self.package_case_metadata(case_info) - print(f"Packaged case metadata for '{self.label}' into {self.destination_dir}") + self.package_case_start() if self.upload: print("Uploading case metadata to the profiling-data repo ...") _push_profiling_data(self.output_root, self.run_commit) # Then package and push each run's results as soon as its own job finishes, # rather than waiting for every run to complete first. - num_runs = len(self.job_infos) - print(f"Waiting for {num_runs} run(s) of '{self.label}' to complete...") - packaged_launch_ids = [] + print(f"Waiting for {len(self.job_infos)} run(s) of '{self.label}' to complete...") for job_info in self._iter_finished_runs(poll_interval=poll_interval): - # The launch id identifies the run; the rank count is just context. - run = f"run {job_info['launch_id']:02d} ({job_info['ranks']} MPI ranks)" - if not self.package_run(job_info, verbose=True): - print(f"No profiling output for '{self.label}' {run}; nothing to package.") - continue - packaged_launch_ids.append(job_info["launch_id"]) - # Refresh `case_metadata.json` so it lists the runs packaged so far. - self.package_case_metadata(case_info) - print(f"Packaged '{self.label}' {run} into {self.destination_dir}") - if self.upload: - print(f"Uploading '{self.label}' {run} to the profiling-data repo ...") + if self.package_finished_run(job_info) and self.upload: + print(f"Uploading '{self.label}' run {job_info['launch_id']:02d} to the profiling-data repo ...") _push_profiling_data(self.output_root, self.run_commit) + self.report_packaged() + + def package_case_start(self) -> None: + """Write `profiling_case_info.json` and package the case-level metadata (the first step of packaging).""" + print( + f"Writing metadata for '{self.label}' to {self.case_output_root / 'profiling_case_info.json'}", + ) + self._case_info = self.case_info_dict() + (self.case_output_root / "profiling_case_info.json").write_text( + json.dumps(self._case_info, indent=2), + encoding="utf-8", + ) + # The case-level metadata goes in before any run, so the packaged folder is already in + # place when the per-run packaging starts updating it. + self.package_case_metadata(self._case_info) + self._packaged_launch_ids = [] + print(f"Packaged case metadata for '{self.label}' into {self.destination_dir}") + + def package_finished_run(self, job_info: dict) -> bool: + """Package one finished run and refresh the case metadata; whether there was output to package.""" + # The launch id identifies the run; the rank count is just context. + run = f"run {job_info['launch_id']:02d} ({job_info['ranks']} MPI ranks)" + if not self.package_run(job_info, verbose=True): + print(f"No profiling output for '{self.label}' {run}; nothing to package.") + return False + self._packaged_launch_ids.append(job_info["launch_id"]) + # Refresh `case_metadata.json` so it lists the runs packaged so far. + self.package_case_metadata(self._case_info) + print(f"Packaged '{self.label}' {run} into {self.destination_dir}") + return True + + def report_packaged(self) -> None: + """Print what was packaged for this case (and how to push or plot it if it was not uploaded).""" + packaged_launch_ids = self._packaged_launch_ids if packaged_launch_ids: print(f"Packaged {len(packaged_launch_ids)} run(s) of '{self.label}' (ids {sorted(packaged_launch_ids)}):") print(f" - {self.destination_dir}") @@ -712,3 +742,97 @@ def finalize_run(self, poll_interval: float = 10.0) -> None: ) else: print(f"No profiling output found for '{self.label}'; nothing to package.") + + +# --------------------------------------------------------------------------- +# Sweeps: many cases in one process (see sweep.py) +# --------------------------------------------------------------------------- +# Cases whose `finalize_run` was deferred, or None outside a sweep. +_SWEEP_CASES: list | None = None +# Packaging folders prepared by `ProfilingCase.setup_run` in this process. +_PREPARED_OUTPUT_ROOTS: set = set() + + +def begin_sweep() -> None: + """Start collecting cases: from now on `ProfilingCase.finalize_run` only registers its case. + + Several drivers running at the same time would each prepare (delete or re-clone) the + shared packaging folder `_profiling_jobs/profiling-data` and push to it, so a sweep runs + all cases in one process instead: every case submits its jobs right away (SLURM runs them + concurrently), and `finalize_all` packages them one at a time as they finish. + """ + global _SWEEP_CASES + _SWEEP_CASES = [] + + +def end_sweep() -> list: + """Stop collecting cases; returns the cases registered since `begin_sweep`.""" + global _SWEEP_CASES + cases, _SWEEP_CASES = _SWEEP_CASES or [], None + return cases + + +def _iter_finished_jobs(cases: list, poll_interval: float = 10.0): + """Yield `(case, job_info)` for every run of `cases` as soon as it finishes, in completion order. + + One poll of the SLURM queue serves all cases. Runs without a SLURM job (local runs, which + `launch` already ran to completion) are yielded first, in launch order. + """ + pending = {} + for case in cases: + for info in case.job_infos: + if info.get("job_id") is None: + yield case, info + else: + pending[info["job_id"]] = (case, info) + + queue = SQueue(user=getpass.getuser()) if pending else None + while pending: + queue.refresh() + active = {job.job_id for job in queue.jobs() if job.is_active} + finished = [job_id for job_id in pending if job_id not in active] + states = job_states(finished) + for job_id in finished: + case, info = pending.pop(job_id) + state = states.get(job_id) + if state is not None and state != "COMPLETED": + print(f"WARNING: job {job_id} for '{case.label}' ({info['ranks']} MPI ranks) ended in state {state}.") + else: + print(f"Job {job_id} for '{case.label}' ({info['ranks']} MPI ranks) finished.") + yield case, info + if pending: + n_cases = len({id(case) for case, _ in pending.values()}) + print(f"{len(pending)} run(s) of {n_cases} case(s) still pending ...", flush=True) + time.sleep(poll_interval) + + +def finalize_all(cases: list, poll_interval: float = 30.0) -> None: + """`ProfilingCase.finalize_run` for several cases at once: package (and push) every run as it finishes. + + The case metadata of all cases is packaged first (one push), then every finished run is + packaged into its case and pushed, one at a time, in the order the jobs finish. + + Args: + cases: The cases, e.g. from `end_sweep`. With `upload`, all must share one packaging + folder (`output_root`), as cases of one process do. + poll_interval: Seconds between polls of the SLURM queue. + """ + if not cases: + print("No cases to finalize.") + return + for case in cases: + case.package_case_start() + uploads = [case for case in cases if case.upload] + if uploads: + print("Uploading the case metadata to the profiling-data repo ...") + _push_profiling_data(uploads[0].output_root, uploads[0].run_commit) + + n_runs = sum(len(case.job_infos) for case in cases) + print(f"Waiting for {n_runs} run(s) of {len(cases)} case(s) to complete ...") + for case, info in _iter_finished_jobs(cases, poll_interval=poll_interval): + if case.package_finished_run(info) and case.upload: + print(f"Uploading '{case.label}' run {info['launch_id']:02d} to the profiling-data repo ...") + _push_profiling_data(case.output_root, case.run_commit) + + for case in cases: + case.report_packaged() diff --git a/profiling/sweep.py b/profiling/sweep.py new file mode 100644 index 000000000..24af87f94 --- /dev/null +++ b/profiling/sweep.py @@ -0,0 +1,121 @@ +"""Run several profiling submit scripts as one sweep, in one process. + +Each submit script (e.g. ``examples/Poisson/gaussian_cube/submit_gaussian_cube.py``) builds a +:class:`profiling_job.ProfilingCase`, launches its runs and calls ``finalize_run``. Started as separate processes, +the scripts would each prepare (delete or re-clone) the shared packaging folder and push to it. A sweep runs them all +in this process between :func:`profiling_job.begin_sweep` and :func:`profiling_job.end_sweep`: every script submits +its jobs right away (SLURM runs them as the partitions allow) and only registers its case, and +:func:`profiling_job.finalize_all` then packages (and with ``upload`` pushes) every run as its job finishes. + +Usage, e.g. ``examples/Poisson/sweep_poisson.py``:: + + from sweep import run_sweep, variant_entries + + entries = variant_entries(Path("examples/Poisson/gaussian_cube/submit_gaussian_cube.py")) + n_failed = run_sweep(entries, upload=False) +""" + +import importlib.util +import runpy +import shlex +import sys +import traceback +from pathlib import Path + +import profiling_job + + +def _load(script: Path): + """Import a submit script as a module (its ``main`` does not run).""" + name = f"_sweep_{script.stem}_{abs(hash(str(script)))}" + spec = importlib.util.spec_from_file_location(name, script) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def variant_entries(script) -> list: + """The entries of a submit script for :func:`run_sweep`: one per key of its ``VARIANTS`` dict, with + ``--variant ``; a script without ``VARIANTS`` gives one entry without arguments. + + Parameters + ---------- + script : str | Path + Path of the submit script. + + Returns + ------- + list[tuple[Path, list[str]]] + ``(script, arguments)`` per run of the script. + """ + script = Path(script).resolve() + variants = getattr(_load(script), "VARIANTS", None) + if not variants: + return [(script, [])] + return [(script, ["--variant", key]) for key in variants] + + +def _command(script: Path, args: list) -> str: + return " ".join(shlex.quote(part) for part in ["python", str(script), *args]) + + +def run_sweep(entries, upload: bool = False, dry_run: bool = False, poll_interval: float = 30.0) -> int: + """Run the submit scripts of ``entries`` in this process, then package (and push) all runs as they finish. + + Every script runs as ``__main__`` with its arguments (plus ``--upload``), inside + :func:`profiling_job.begin_sweep`, so its ``finalize_run`` only registers its case; a script that fails is + reported and the sweep goes on. :func:`profiling_job.finalize_all` then waits on the runs of all registered cases + together. + + Parameters + ---------- + entries : list[tuple[str | Path, list[str]]] + ``(script, arguments)`` per submit command, e.g. from :func:`variant_entries`. + upload : bool + Pass ``--upload`` to every script (push the packaged runs to the profiling-data repo). + dry_run : bool + Only print the commands. + poll_interval : float + Seconds between polls of the SLURM queue in :func:`profiling_job.finalize_all`. + + Returns + ------- + int + Number of scripts that failed. + """ + entries = [(Path(script).resolve(), list(args) + (["--upload"] if upload else [])) for script, args in entries] + print(f"Sweep of {len(entries)} submit command(s):") + for script, args in entries: + print(f" {_command(script, args)}") + if dry_run: + return 0 + + failed = [] + profiling_job.begin_sweep() + try: + for script, args in entries: + print(f"\n=== {_command(script, args)}", flush=True) + argv = sys.argv + sys.argv = [str(script), *args] + try: + runpy.run_path(str(script), run_name="__main__") + except SystemExit as exc: + if exc.code not in (None, 0): + failed.append((script, args)) + print(f"FAILED (exit code {exc.code}): {_command(script, args)}") + except Exception: + failed.append((script, args)) + traceback.print_exc() + print(f"FAILED: {_command(script, args)}") + finally: + sys.argv = argv + finally: + cases = profiling_job.end_sweep() + + profiling_job.finalize_all(cases, poll_interval=poll_interval) + + if failed: + print(f"\n{len(failed)} of {len(entries)} submit command(s) failed:") + for script, args in failed: + print(f" {_command(script, args)}") + return len(failed)