diff --git a/profiling/clusters.py b/profiling/clusters.py index d64c9a1f6..483ff83b4 100644 --- a/profiling/clusters.py +++ b/profiling/clusters.py @@ -69,6 +69,19 @@ def detect_machine_name() -> str | None: "mail_type": "none", "time": "00:15:00", }, + # Pitagora's GPU (booster) partition, one node with its 4 GPUs. Not detected by `detect_machine_name` + # (the login nodes are shared with the DCGP partition), so a run picks it with `launch(slurm_preset=...)`. + "pitagora_booster": { + "cpus_per_task": 1, + "mem": "480GB", + "gres": "gpu:4,tmpfs:10g", + "partition": "boost_fua_dbg", + "account": "FUSIO_HLST_6", + "output": "./%x.%j.out", + "error": "./%x.%j.err", + "mail_type": "none", + "time": "00:15:00", + }, "tok": { "cpus_per_task": 1, "mem_per_cpu": "1GB", diff --git a/profiling/examples/VlasovAmpereOneSpecies/landau_damping/params_vlasov_ampere.py b/profiling/examples/VlasovAmpereOneSpecies/landau_damping/params_vlasov_ampere.py new file mode 100644 index 000000000..76d3ae9d0 --- /dev/null +++ b/profiling/examples/VlasovAmpereOneSpecies/landau_damping/params_vlasov_ampere.py @@ -0,0 +1,239 @@ +# ----------------------------- +# 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-Ampere weak Landau damping" +description = """ +Weak Landau damping with the VlasovAmpereOneSpecies model, the setup of the verification test in +src/struphy/models/tests/verification: periodic box of length 4 pi (k = 0.5), Maxwellian electrons with a density +perturbation of 1e-3 cos(k x), control-variate weights, cubic splines in x. The electric field energy decays at the +rate -0.1533. alpha = 1 and epsilon = -1 (electrons, time in units of the inverse plasma frequency, velocities in +units of c). + +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: the run activates +the profiling session of Simulation.run, and on the CuPy backend the session records CUDA-event device time per +region (GPUOptions.timing) and synchronizes the device at region ends, so the CPU time of a region includes its +kernels. The post-processing writes the scope-profiler HTML report and a summary of the cost per step and marker +into the run's results folder. +""" + +import argparse +import logging + +from struphy import set_logging_level + +set_logging_level(logging.INFO) + +import math + +import cunumpy as xp + +# ------------------ +# Import Struphy API +# ------------------ +from struphy import ( + BoundaryParameters, + DerhamOptions, + EnvironmentOptions, + LoadingParameters, + ProfilingOptions, + SavingParameters, + Simulation, + SortingParameters, + Time, + WeightsParameters, + domains, + grids, + maxwellians, + perturbations, +) + +# -------------------------- +# Command line parameters +# -------------------------- +# `--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.") +args, _ = parser.parse_known_args() + +# --------------------- +# Instance of the model +# --------------------- +from struphy.models import VlasovAmpereOneSpecies + +model = VlasovAmpereOneSpecies(alpha=1.0, epsilon=-1.0, with_B0=False) + +# List all variables and decide whether to save their data +model.em_fields.e_field.save_data = True +model.kinetic_ions.var.save_data = False + +# -------------------------- +# Instance of the simulation +# -------------------------- + +# Environment options +env = EnvironmentOptions( + sim_folder=f"sim_{args.id:02d}", + restart=False, + save_restart=False, # restart data would dominate the run time for many markers +) + +# Time stepping +dt = 0.05 +steps = 20 +time_opts = Time(dt=dt, Tend=round(steps * dt, 14)) + +# Geometry +Lx = 4 * math.pi # k = 0.5 +domain = domains.Cuboid(r1=Lx) + +# Fluid equilibrium (can be used as part of initial conditions) +equil = None + +# Grid +# cells in x (1 in y and z). Only x is decomposed, and every rank needs at least `degree` = 3 cells of its own, +# so 1024 cells run on up to 256 ranks (4 cells per rank), one full CPU node. +num_elements = 1024 +grid = grids.TensorProductGrid(num_elements=(num_elements, 1, 1), mpi_dims_mask=(True, False, False)) + +# Derham options +derham_opts = DerhamOptions(degree=(3, 1, 1)) + +# 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 comparison plots of the profiling-data page: the runs of +# one case differ in backend and rank count only, so the label names those two (e.g. "cupy, 4 ranks"). +from maybempi import MPI, SerialMPI + +num_ranks = 1 if isinstance(MPI, SerialMPI) else MPI.COMM_WORLD.Get_size() +profiling_label = f"{xp.get_backend()}, {num_ranks} rank" + ("s" if num_ranks != 1 else "") + +profiling_opts = ProfilingOptions( + label=profiling_label, + use_gpu_timing=gpu_timing, + gpu_timing_backend="cupy", + gpu_sync_on_exit=gpu_timing, +) + +# 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 +# ------------------- + +ppc = 1000 +loading_params = LoadingParameters(ppc=ppc, seed=1234) +weights_params = WeightsParameters(control_variate=True) +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 physics, 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, + bufsize=0.4, +) + +# ------------------ +# Propagator options +# ------------------ + +for prop in model.propagators.__dict__.values(): + prop.options = prop.Options() +model.initial_poisson.options = model.initial_poisson.Options(stab_mat="M0") + +# ------------------ +# Initial conditions +# ------------------ +# Initial conditions are the sum of the background(s) and the perturbation(s). +# For kinetic species the background is mandatory; the perturbation is added to the density. + +background = maxwellians.Maxwellian3D(n=(1.0, None)) +model.kinetic_ions.var.add_background(background) + +perturbation = perturbations.ModesCos(ls=(1,), amps=(1e-3,)) +model.kinetic_ions.var.add_initial_condition(maxwellians.Maxwellian3D(n=(1.0, perturbation))) + + +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 + + with device_context() if fake_cupy_active() else contextlib.nullcontext(): + out = sim.run(profiling_activated=True) + + if sim.rank == 0: + # `out.path_out` is the run's (absolute) output folder. The profiling packaging picks the files + # written to `results` up from here and uploads them next to the run's profiling_data.h5. + import json + import statistics + + import numpy as np + from scope_profiler import create_html_report + + 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) + steps = results["model.integrate"] + per_rank = [steps[rank].durations for rank in steps.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)) + n_markers = ppc * num_elements + summary = { + "backend": xp.get_backend(), + "gpu_timing": gpu_timing, + "mpi_ranks": results.num_ranks, + "num_elements": num_elements, + "ppc": ppc, + "n_markers": n_markers, + "steps": steps, + "dt": dt, + "step_times": step_times.tolist(), + "step_median": step_median, + "ns_per_marker_step": 1e9 * step_median / 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 ({n_markers:,} markers)" + ) + # the propagators and kernels of the run; on the CuPy backend with GPU columns + results.print_summary(include=["^prop: ", "^kernel: "], sort="total", title="propagators and kernels") diff --git a/profiling/examples/VlasovAmpereOneSpecies/submit_vlasov_ampere_landau_damping.py b/profiling/examples/VlasovAmpereOneSpecies/submit_vlasov_ampere_landau_damping.py new file mode 100644 index 000000000..86a3cb378 --- /dev/null +++ b/profiling/examples/VlasovAmpereOneSpecies/submit_vlasov_ampere_landau_damping.py @@ -0,0 +1,81 @@ +"""Vlasov-Ampere weak Landau damping profiling case. + +This file defines the Vlasov-Ampere weak Landau damping profiling case (the `ProfilingCase`) and submits +it: for each rank count, `ProfilingCase.launch` builds and submits a SLURM script +(using `clusters.SLURM_PRESETS` by default), or, without a batch system, runs +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_vlasov_ampere.py` +directly (its `__main__` block is the worker). + +The run is profiled with scope-profiler. `--backend cupy` runs it on the GPU with CUDA-event +timing per region; the job script exports the choice as CUNUMPY_BACKEND before the simulation. +The resolution (1024 cells, 1000 markers per cell, 20 steps) is set in `params_vlasov_ampere.py`. +""" + +import argparse +import sys +from pathlib import Path + +# the profiling driver (profiling_job.py, clusters.py, ...) lives in profiling/, two levels up +sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + +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( + "--backend", + default="numpy", + choices=["numpy", "cupy"], + help="cunumpy backend to run on; cupy records CUDA-event GPU times per region.", + ) + args = parser.parse_args() + + # Paths relative to this script's location, so it can be run from anywhere. + script_dir = Path(__file__).resolve().parent + params_dir = script_dir / "landau_damping" + + profiling_case = ProfilingCase( + label="vlasov_ampere_landau_damping", + name="Vlasov-Ampere weak Landau damping", + description="Strong scaling of the VlasovAmpereOneSpecies model on weak Landau damping (periodic box, k = 0.5, 1e-3 density perturbation), on the NumPy or CuPy backend.", + physics_problem="Landau damping of an electrostatic wave in a collisionless electron plasma.", + struphy_model_used="VlasovAmpereOneSpecies", + params_source=params_dir / "params_vlasov_ampere.py", + language=args.language, + compiler=args.compiler, + upload=args.upload, + ) + + # Launch one run per rank count + for num_tasks in (1, 2, 4, 8, 16, 32, 64, 128, 256): + profiling_case.launch(num_tasks, backend=args.backend) + + # Package and push each run as its own job finishes. + profiling_case.finalize_run() + + +if __name__ == "__main__": + main() diff --git a/profiling/examples/VlasovAmpereOneSpecies/submit_vlasov_ampere_landau_damping_cpu_node_vs_gpu_node.py b/profiling/examples/VlasovAmpereOneSpecies/submit_vlasov_ampere_landau_damping_cpu_node_vs_gpu_node.py new file mode 100644 index 000000000..3cd626e76 --- /dev/null +++ b/profiling/examples/VlasovAmpereOneSpecies/submit_vlasov_ampere_landau_damping_cpu_node_vs_gpu_node.py @@ -0,0 +1,77 @@ +"""Vlasov-Ampere weak Landau damping profiling case: one CPU node vs one GPU node. + +This file defines the profiling case (the `ProfilingCase`) and submits its two runs of the same +problem: one full CPU node (256 MPI ranks, NumPy backend) and one full GPU node (4 MPI ranks, one per +GPU, CuPy backend). The two runs go to different SLURM partitions, so each `ProfilingCase.launch` +names its own preset from `clusters.SLURM_PRESETS`. Without a batch system the runs execute directly +on this machine, where the 256-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_vlasov_ampere.py` +directly (its `__main__` block is the worker). + +The runs are profiled with scope-profiler, the CuPy run with CUDA-event timing per region. The +resolution (1024 cells, 1000 markers per cell, 20 steps) is set in `params_vlasov_ampere.py`, shared +with the strong-scaling case of `submit_vlasov_ampere_landau_damping.py`. +""" + +import argparse +import sys +from pathlib import Path + +# the profiling driver (profiling_job.py, clusters.py, ...) lives in profiling/, two levels up +sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + +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 + params_dir = script_dir / "landau_damping" + + profiling_case = ProfilingCase( + label="vlasov_ampere_landau_damping_cpu_node_vs_gpu_node", + name="Vlasov-Ampere weak Landau damping: CPU node vs GPU node", + description="The VlasovAmpereOneSpecies model on weak Landau damping (periodic box, k = 0.5, 1e-3 density perturbation, 1024 cells, about 1M markers) on one full CPU node (256 MPI ranks, NumPy) vs one full GPU node (4 MPI ranks on 4 GPUs, CuPy).", + physics_problem="Landau damping of an electrostatic wave in a collisionless electron plasma.", + struphy_model_used="VlasovAmpereOneSpecies", + params_source=params_dir / "params_vlasov_ampere.py", + language=args.language, + compiler=args.compiler, + upload=args.upload, + ) + + # One CPU node: 256 ranks on Pitagora's DCGP partition + profiling_case.launch(256, num_nodes=1, backend="numpy", slurm_preset="pitagora_dcgp") + # One GPU node: one rank per GPU on Pitagora's booster partition + profiling_case.launch(4, 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/profiling/examples/VlasovMaxwellOneSpecies/landau_damping/params_vlasov_maxwell.py b/profiling/examples/VlasovMaxwellOneSpecies/landau_damping/params_vlasov_maxwell.py new file mode 100644 index 000000000..fa307566d --- /dev/null +++ b/profiling/examples/VlasovMaxwellOneSpecies/landau_damping/params_vlasov_maxwell.py @@ -0,0 +1,238 @@ +# ----------------------------- +# 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-Maxwell weak Landau damping" +description = """ +Weak Landau damping with the VlasovMaxwellOneSpecies model, the setup of the verification test in +src/struphy/models/tests/verification: periodic box of length 4 pi (k = 0.5), Maxwellian electrons with a density +perturbation of 1e-3 cos(k x), control-variate weights, cubic splines in x. The electric field energy decays at the +rate -0.1533. alpha = 1 and epsilon = -1 (electrons, time in units of the inverse plasma frequency, velocities in +units of c). + +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: the run activates +the profiling session of Simulation.run, and on the CuPy backend the session records CUDA-event device time per +region (GPUOptions.timing) and synchronizes the device at region ends, so the CPU time of a region includes its +kernels. The post-processing writes the scope-profiler HTML report and a summary of the cost per step and marker +into the run's results folder. +""" + +import argparse +import logging + +from struphy import set_logging_level + +set_logging_level(logging.INFO) + +import math + +import cunumpy as xp + +# ------------------ +# Import Struphy API +# ------------------ +from struphy import ( + BoundaryParameters, + DerhamOptions, + EnvironmentOptions, + LoadingParameters, + ProfilingOptions, + SavingParameters, + Simulation, + SortingParameters, + Time, + WeightsParameters, + domains, + grids, + maxwellians, + perturbations, +) + +# -------------------------- +# Command line parameters +# -------------------------- +# `--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.") +args, _ = parser.parse_known_args() + +# --------------------- +# Instance of the model +# --------------------- +from struphy.models import VlasovMaxwellOneSpecies + +model = VlasovMaxwellOneSpecies(alpha=1.0, epsilon=-1.0) + +# List all variables and decide whether to save their data +model.em_fields.e_field.save_data = True +model.em_fields.b_field.save_data = True +model.kinetic_ions.var.save_data = False + +# -------------------------- +# Instance of the simulation +# -------------------------- + +# Environment options +env = EnvironmentOptions( + sim_folder=f"sim_{args.id:02d}", + restart=False, + save_restart=False, # restart data would dominate the run time for many markers +) + +# Time stepping +dt = 0.05 +steps = 20 +time_opts = Time(dt=dt, Tend=round(steps * dt, 14)) + +# Geometry +Lx = 4 * math.pi # k = 0.5 +domain = domains.Cuboid(r1=Lx) + +# Fluid equilibrium (can be used as part of initial conditions) +equil = None + +# Grid +num_elements = 32 # cells in x (1 in y and z) +grid = grids.TensorProductGrid(num_elements=(num_elements, 1, 1)) + +# Derham options +derham_opts = DerhamOptions(degree=(3, 1, 1)) + +# 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 comparison plots of the profiling-data page: the runs of +# one case differ in backend and rank count only, so the label names those two (e.g. "cupy, 4 ranks"). +from maybempi import MPI, SerialMPI + +num_ranks = 1 if isinstance(MPI, SerialMPI) else MPI.COMM_WORLD.Get_size() +profiling_label = f"{xp.get_backend()}, {num_ranks} rank" + ("s" if num_ranks != 1 else "") + +profiling_opts = ProfilingOptions( + label=profiling_label, + use_gpu_timing=gpu_timing, + gpu_timing_backend="cupy", + gpu_sync_on_exit=gpu_timing, +) + +# 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 +# ------------------- + +ppc = 1000 +loading_params = LoadingParameters(ppc=ppc, seed=1234) +weights_params = WeightsParameters(control_variate=True) +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 physics, 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, + bufsize=0.4, +) + +# ------------------ +# Propagator options +# ------------------ + +for prop in model.propagators.__dict__.values(): + prop.options = prop.Options() +model.initial_poisson.options = model.initial_poisson.Options(stab_mat="M0") + +# ------------------ +# Initial conditions +# ------------------ +# Initial conditions are the sum of the background(s) and the perturbation(s). +# For kinetic species the background is mandatory; the perturbation is added to the density. + +background = maxwellians.Maxwellian3D(n=(1.0, None)) +model.kinetic_ions.var.add_background(background) + +perturbation = perturbations.ModesCos(ls=(1,), amps=(1e-3,)) +model.kinetic_ions.var.add_initial_condition(maxwellians.Maxwellian3D(n=(1.0, perturbation))) + + +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 + + with device_context() if fake_cupy_active() else contextlib.nullcontext(): + out = sim.run(profiling_activated=True) + + if sim.rank == 0: + # `out.path_out` is the run's (absolute) output folder. The profiling packaging picks the files + # written to `results` up from here and uploads them next to the run's profiling_data.h5. + import json + import statistics + + import numpy as np + from scope_profiler import create_html_report + + 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) + steps = results["model.integrate"] + per_rank = [steps[rank].durations for rank in steps.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)) + n_markers = ppc * num_elements + summary = { + "backend": xp.get_backend(), + "gpu_timing": gpu_timing, + "mpi_ranks": results.num_ranks, + "num_elements": num_elements, + "ppc": ppc, + "n_markers": n_markers, + "steps": steps, + "dt": dt, + "step_times": step_times.tolist(), + "step_median": step_median, + "ns_per_marker_step": 1e9 * step_median / 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 ({n_markers:,} markers)" + ) + # the propagators and kernels of the run; on the CuPy backend with GPU columns + results.print_summary(include=["^prop: ", "^kernel: "], sort="total", title="propagators and kernels") diff --git a/profiling/examples/VlasovMaxwellOneSpecies/submit_vlasov_maxwell_landau_damping.py b/profiling/examples/VlasovMaxwellOneSpecies/submit_vlasov_maxwell_landau_damping.py new file mode 100644 index 000000000..131aac72d --- /dev/null +++ b/profiling/examples/VlasovMaxwellOneSpecies/submit_vlasov_maxwell_landau_damping.py @@ -0,0 +1,81 @@ +"""Vlasov-Maxwell weak Landau damping profiling case. + +This file defines the Vlasov-Maxwell weak Landau damping profiling case (the `ProfilingCase`) and submits +it: for each rank count, `ProfilingCase.launch` builds and submits a SLURM script +(using `clusters.SLURM_PRESETS` by default), or, without a batch system, runs +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_vlasov_maxwell.py` +directly (its `__main__` block is the worker). + +The run is profiled with scope-profiler. `--backend cupy` runs it on the GPU with CUDA-event +timing per region; the job script exports the choice as CUNUMPY_BACKEND before the simulation. +The resolution (32 cells, 1000 markers per cell, 20 steps) is set in `params_vlasov_maxwell.py`. +""" + +import argparse +import sys +from pathlib import Path + +# the profiling driver (profiling_job.py, clusters.py, ...) lives in profiling/, two levels up +sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + +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( + "--backend", + default="numpy", + choices=["numpy", "cupy"], + help="cunumpy backend to run on; cupy records CUDA-event GPU times per region.", + ) + args = parser.parse_args() + + # Paths relative to this script's location, so it can be run from anywhere. + script_dir = Path(__file__).resolve().parent + params_dir = script_dir / "landau_damping" + + profiling_case = ProfilingCase( + label="vlasov_maxwell_landau_damping", + name="Vlasov-Maxwell weak Landau damping", + description="Strong scaling of the VlasovMaxwellOneSpecies model on weak Landau damping (periodic box, k = 0.5, 1e-3 density perturbation), on the NumPy or CuPy backend.", + physics_problem="Landau damping of an electrostatic wave in a collisionless electron plasma.", + struphy_model_used="VlasovMaxwellOneSpecies", + params_source=params_dir / "params_vlasov_maxwell.py", + language=args.language, + compiler=args.compiler, + upload=args.upload, + ) + + # Launch one run per rank count + for num_tasks in (1, 2, 4, 8, 16, 32, 64, 128, 256): + profiling_case.launch(num_tasks, backend=args.backend) + + # Package and push each run as its own job finishes. + profiling_case.finalize_run() + + +if __name__ == "__main__": + main() diff --git a/profiling/examples/VlasovMaxwellOneSpecies/submit_vlasov_maxwell_weibel_instability.py b/profiling/examples/VlasovMaxwellOneSpecies/submit_vlasov_maxwell_weibel_instability.py new file mode 100644 index 000000000..2da4c7315 --- /dev/null +++ b/profiling/examples/VlasovMaxwellOneSpecies/submit_vlasov_maxwell_weibel_instability.py @@ -0,0 +1,81 @@ +"""Vlasov-Maxwell Weibel instability profiling case. + +This file defines the Vlasov-Maxwell Weibel instability profiling case (the `ProfilingCase`) and submits +it: for each rank count, `ProfilingCase.launch` builds and submits a SLURM script +(using `clusters.SLURM_PRESETS` by default), or, without a batch system, runs +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_vlasov_maxwell.py` +directly (its `__main__` block is the worker). + +The run is profiled with scope-profiler. `--backend cupy` runs it on the GPU with CUDA-event +timing per region; the job script exports the choice as CUNUMPY_BACKEND before the simulation. +The resolution (32 cells, 1000 markers per cell, 20 steps) is set in `params_vlasov_maxwell.py`. +""" + +import argparse +import sys +from pathlib import Path + +# the profiling driver (profiling_job.py, clusters.py, ...) lives in profiling/, two levels up +sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + +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( + "--backend", + default="numpy", + choices=["numpy", "cupy"], + help="cunumpy backend to run on; cupy records CUDA-event GPU times per region.", + ) + args = parser.parse_args() + + # Paths relative to this script's location, so it can be run from anywhere. + script_dir = Path(__file__).resolve().parent + params_dir = script_dir / "weibel_instability" + + profiling_case = ProfilingCase( + label="vlasov_maxwell_weibel_instability", + name="Vlasov-Maxwell Weibel instability", + description="Strong scaling of the VlasovMaxwellOneSpecies model on the Weibel instability (Kraus et al. 2017, Sec. 6.2: anisotropic Maxwellian, seed field B3 = -1e-4 cos(k x)), on the NumPy or CuPy backend.", + physics_problem="Weibel instability driven by a temperature anisotropy in a collisionless electron plasma.", + struphy_model_used="VlasovMaxwellOneSpecies", + params_source=params_dir / "params_vlasov_maxwell.py", + language=args.language, + compiler=args.compiler, + upload=args.upload, + ) + + # Launch one run per rank count + for num_tasks in (1, 2, 4, 8, 16, 32, 64, 128, 256): + profiling_case.launch(num_tasks, backend=args.backend) + + # Package and push each run as its own job finishes. + profiling_case.finalize_run() + + +if __name__ == "__main__": + main() diff --git a/profiling/examples/VlasovMaxwellOneSpecies/weibel_instability/params_vlasov_maxwell.py b/profiling/examples/VlasovMaxwellOneSpecies/weibel_instability/params_vlasov_maxwell.py new file mode 100644 index 000000000..e325b3c92 --- /dev/null +++ b/profiling/examples/VlasovMaxwellOneSpecies/weibel_instability/params_vlasov_maxwell.py @@ -0,0 +1,243 @@ +# ----------------------------- +# 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-Maxwell Weibel instability" +description = """ +Weibel instability with the VlasovMaxwellOneSpecies model, the setup of Kraus et al., J. Plasma Phys. 83 (2017), +Sec. 6.2, as in the verification test in src/struphy/models/tests/verification: periodic box of length 2 pi / 1.25, +anisotropic Maxwellian electrons (vth1 = vth3 = 0.02 / sqrt(2), vth2 = sqrt(12) vth1), seed field +B3 = -1e-4 cos(k x), no density perturbation, control-variate weights, cubic splines in x. The magnetic field energy +grows at the rate 0.02784. alpha = 1 and epsilon = -1 (electrons, time in units of the inverse plasma frequency, +velocities in units of c). + +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: the run activates +the profiling session of Simulation.run, and on the CuPy backend the session records CUDA-event device time per +region (GPUOptions.timing) and synchronizes the device at region ends, so the CPU time of a region includes its +kernels. The post-processing writes the scope-profiler HTML report and a summary of the cost per step and marker +into the run's results folder. +""" + +import argparse +import logging + +from struphy import set_logging_level + +set_logging_level(logging.INFO) + +import math + +import cunumpy as xp + +# ------------------ +# Import Struphy API +# ------------------ +from struphy import ( + BoundaryParameters, + DerhamOptions, + EnvironmentOptions, + LoadingParameters, + ProfilingOptions, + SavingParameters, + Simulation, + SortingParameters, + Time, + WeightsParameters, + domains, + grids, + maxwellians, + perturbations, +) + +# -------------------------- +# Command line parameters +# -------------------------- +# `--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.") +args, _ = parser.parse_known_args() + +# --------------------- +# Instance of the model +# --------------------- +from struphy.models import VlasovMaxwellOneSpecies + +model = VlasovMaxwellOneSpecies(alpha=1.0, epsilon=-1.0) + +# List all variables and decide whether to save their data +model.em_fields.e_field.save_data = True +model.em_fields.b_field.save_data = True +model.kinetic_ions.var.save_data = False + +# -------------------------- +# Instance of the simulation +# -------------------------- + +# Environment options +env = EnvironmentOptions( + sim_folder=f"sim_{args.id:02d}", + restart=False, + save_restart=False, # restart data would dominate the run time for many markers +) + +# Time stepping +dt = 0.05 +steps = 20 +time_opts = Time(dt=dt, Tend=round(steps * dt, 14)) + +# Geometry +Lx = 2 * math.pi / 1.25 # k = 1.25 +domain = domains.Cuboid(r1=Lx) + +# Fluid equilibrium (can be used as part of initial conditions) +equil = None + +# Grid +num_elements = 32 # cells in x (1 in y and z) +grid = grids.TensorProductGrid(num_elements=(num_elements, 1, 1)) + +# Derham options +derham_opts = DerhamOptions(degree=(3, 1, 1)) + +# 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 comparison plots of the profiling-data page: the runs of +# one case differ in backend and rank count only, so the label names those two (e.g. "cupy, 4 ranks"). +from maybempi import MPI, SerialMPI + +num_ranks = 1 if isinstance(MPI, SerialMPI) else MPI.COMM_WORLD.Get_size() +profiling_label = f"{xp.get_backend()}, {num_ranks} rank" + ("s" if num_ranks != 1 else "") + +profiling_opts = ProfilingOptions( + label=profiling_label, + use_gpu_timing=gpu_timing, + gpu_timing_backend="cupy", + gpu_sync_on_exit=gpu_timing, +) + +# 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 +# ------------------- + +ppc = 1000 +loading_params = LoadingParameters(ppc=ppc, seed=1234) +weights_params = WeightsParameters(control_variate=True) +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 physics, 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, + bufsize=0.4, +) + +# ------------------ +# Propagator options +# ------------------ + +for prop in model.propagators.__dict__.values(): + prop.options = prop.Options() +model.initial_poisson.options = model.initial_poisson.Options(stab_mat="M0") + +# ------------------ +# Initial conditions +# ------------------ +# Initial conditions are the sum of the background(s) and the perturbation(s). +# For kinetic species the background is mandatory; without add_initial_condition() it is the initial condition. + +vth1 = 0.02 / math.sqrt(2.0) +background = maxwellians.Maxwellian3D( + n=(1.0, None), vth1=(vth1, None), vth2=(math.sqrt(12.0) * vth1, None), vth3=(vth1, None) +) +model.kinetic_ions.var.add_background(background) + +# B3 = -1e-4 cos(k x) in physical coordinates: kx = 2 pi l / Lx, so Lx must be the box length +seed = perturbations.ModesCos(ls=(1,), amps=(-1e-4,), Lx=Lx, given_in_basis="physical", comp=2) +model.em_fields.b_field.add_perturbation(seed) + + +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 + + with device_context() if fake_cupy_active() else contextlib.nullcontext(): + out = sim.run(profiling_activated=True) + + if sim.rank == 0: + # `out.path_out` is the run's (absolute) output folder. The profiling packaging picks the files + # written to `results` up from here and uploads them next to the run's profiling_data.h5. + import json + import statistics + + import numpy as np + from scope_profiler import create_html_report + + 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) + steps = results["model.integrate"] + per_rank = [steps[rank].durations for rank in steps.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)) + n_markers = ppc * num_elements + summary = { + "backend": xp.get_backend(), + "gpu_timing": gpu_timing, + "mpi_ranks": results.num_ranks, + "num_elements": num_elements, + "ppc": ppc, + "n_markers": n_markers, + "steps": steps, + "dt": dt, + "step_times": step_times.tolist(), + "step_median": step_median, + "ns_per_marker_step": 1e9 * step_median / 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 ({n_markers:,} markers)" + ) + # the propagators and kernels of the run; on the CuPy backend with GPU columns + results.print_summary(include=["^prop: ", "^kernel: "], sort="total", title="propagators and kernels") diff --git a/profiling/package_profiling_results.py b/profiling/package_profiling_results.py index 2512f1390..180bc17b3 100644 --- a/profiling/package_profiling_results.py +++ b/profiling/package_profiling_results.py @@ -13,9 +13,10 @@ RUN_METADATA_FILE = "run_metadata.json" # Written by the parameter file itself (see the Poisson example), one per `sim_` run -# directory: post-processing figures and small arrays, next to the profiling output. +# directory: post-processing figures, small arrays, scope-profiler HTML reports and JSON summaries +# (see the Vlasov examples), next to the profiling output. RESULTS_DIR_NAME = "results" -RESULTS_FILE_SUFFIXES = (".png", ".npy") +RESULTS_FILE_SUFFIXES = (".png", ".npy", ".html", ".json") def _run_name(launch_id: int) -> str: @@ -85,6 +86,8 @@ def _write_run_metadata( # Struphy records the rank count it actually ran with; fall back to the requested one # for a run that wrote no metadata of its own. metadata.setdefault("mpi_ranks", job_info.get("ranks")) + # The cunumpy backend is chosen per launch, so it is recorded per run. + metadata["backend"] = job_info.get("backend") metadata["slurm_script"] = job_info.get("slurm_dict") metadata["slurm_script_str"] = job_info.get("job_script") metadata["packaged_files"] = { diff --git a/profiling/profiling_job.py b/profiling/profiling_job.py index 10a7293db..d2a37b8f4 100644 --- a/profiling/profiling_job.py +++ b/profiling/profiling_job.py @@ -7,13 +7,14 @@ results directories. It stores everything it computes as attributes on the case itself (`venv_path`, `use_slurm`, `launcher`, `case_output_root`, ...), so there's a single object to thread through the rest of the run. -- For each rank count, the caller calls `ProfilingCase.launch(ntasks, param_flags=None)`, which builds +- 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. - Under SLURM, the cluster preset is picked on each call from - `clusters.SLURM_PRESETS` unless a `slurm_presets` argument overrides it. + 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. - `ProfilingCase.finalize_run` packages and pushes the case-level metadata straight away, then waits on the runs one by one, packaging and pushing each run's results into that same folder as soon as its own job finishes. With `upload=True`, the @@ -132,7 +133,12 @@ class ProfilingCase: def __post_init__(self) -> None: self.setup_run() - def build_commands(self, ntasks: int, param_flags: list[str] | None = None) -> list[str]: + def build_commands( + self, + ntasks: int, + param_flags: list[str] | None = None, + backend: str = "numpy", + ) -> list[str]: """Build the shell commands that run this case with a single MPI rank count. Reads the setup computed by `setup_run` (`case_output_root`, `venv_path`, @@ -144,6 +150,8 @@ def build_commands(self, ntasks: int, param_flags: list[str] | None = None) -> l param_flags: Extra CLI flags appended to the `params_source` invocation, e.g. `["--ppc", "10"]`. Omit for none. `--id` is always passed on top of these. + backend: cunumpy backend of the run, "numpy" (CPU) or "cupy" (GPU), + exported as CUNUMPY_BACKEND before the simulation. Returns: The shell commands to run, in order, as one script. @@ -171,11 +179,14 @@ def build_commands(self, ntasks: int, param_flags: list[str] | None = None) -> l ), "set -e", f"source {activate_path!s}", + # cunumpy reads its backend from this variable when it is imported, so it is exported before the run + f"export CUNUMPY_BACKEND={backend}", 'echo "----------------------------------------"', f'echo "Running profiling case: {self.label} ({ntasks} MPI ranks)"', f'echo "Description: {self.description}"', f'echo "Physics problem: {self.physics_problem}"', f'echo "Struphy model used: {self.struphy_model_used}"', + f'echo "cunumpy backend: {backend}"', f'echo "Case directory: {output_root}"', 'echo "----------------------------------------"', f'mkdir -p "{sim_dir}"', @@ -202,7 +213,9 @@ def launch( num_tasks: int = 1, num_nodes: int | None = None, param_flags: list[str] | None = None, - slurm_presets: dict[str, dict] | None = None, + backend: str = "numpy", + slurm_preset: str | None = None, + slurm_options: dict | None = None, ) -> None: """Build, submit/run, and record the run for a single rank count. @@ -223,14 +236,24 @@ def launch( outside SLURM, where every rank runs in a single local process group. param_flags: Extra CLI flags forwarded to `build_commands` and appended to the `params_source` invocation, e.g. `["--ppc", "10"]`. - slurm_presets: Candidate SLURM presets, keyed by cluster name; one is - picked via `detect_machine_name` on every call. Defaults to - `clusters.SLURM_PRESETS`. Ignored outside SLURM. + backend: cunumpy backend of this run, "numpy" (CPU) or "cupy" (GPU): + exported as CUNUMPY_BACKEND in the job script and recorded in the + run's metadata, so one case can mix runs on both backends. + slurm_preset: Key of the `clusters.SLURM_PRESETS` entry the job is + submitted with, e.g. "pitagora_booster" for a GPU partition. + Defaults to the cluster found by `detect_machine_name`. Its + `HARDWARE_INFO` entry is only needed when `num_nodes` is not given. + Ignored outside SLURM. + slurm_options: SLURM pragmas that override those of the preset, e.g. + `{"time": "00:30:00"}`. Ignored outside SLURM. Raises: - ValueError: If `num_tasks` is not evenly divisible by `num_nodes` - (SLURM only). + ValueError: If `backend` is neither "numpy" nor "cupy", or if + `num_tasks` is not evenly divisible by `num_nodes` (SLURM only). """ + if backend not in ("numpy", "cupy"): + raise ValueError(f"backend must be 'numpy' or 'cupy', got {backend!r}") + if not self.use_slurm: available = os.cpu_count() or 1 if num_tasks > available: @@ -243,23 +266,19 @@ def launch( script_path = profiling_results_base / f"job_profile_{self.label}_{self.launch_count:02d}.sh" # Build the commands to run this case with the given rank count, and write/submit - case_commands = self.build_commands(num_tasks, param_flags) + case_commands = self.build_commands(num_tasks, param_flags, backend) # Submit a SLURM script or run a local bash script, depending on whether SLURM is available. if self.use_slurm: - cluster_name = detect_machine_name() + # The preset of this run: the caller's choice, or the one of the detected cluster. + preset_name = slurm_preset or detect_machine_name() if num_nodes is None: # Compute the number of nodes needed to run `num_tasks` ranks, given the cluster's hardware. - cpus_per_node = HARDWARE_INFO[cluster_name]["cpus_per_node"] + cpus_per_node = HARDWARE_INFO[preset_name]["cpus_per_node"] num_nodes = (num_tasks + cpus_per_node - 1) // cpus_per_node if num_tasks % num_nodes != 0: raise ValueError(f"num_tasks ({num_tasks}) is not evenly divisible by num_nodes ({num_nodes}).") - - # 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] - else: - cluster_preset = SLURM_PRESETS[cluster_name] + cluster_preset = {**SLURM_PRESETS[preset_name], **(slurm_options or {})} # Build the slurm script script = SlurmScript( @@ -299,6 +318,7 @@ def launch( job_info = { "launch_id": self.launch_count, "ranks": num_tasks, + "backend": backend, "job_id": job_id, "job_script_path": str(script_path), "job_script": script_text,