Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 18 additions & 8 deletions experiments/gromacs/experiment.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,14 +19,14 @@ class Gromacs(
):
variant(
"workload",
default="water_gmx50_adac",
default="water_gmx50",
description="workload name",
)

variant(
"version",
default="2024",
values=("2024", "2023.3"),
default="2025.2",
values=("2025.2", "2024", "2023.3"),
description="app version",
)

Expand All @@ -36,12 +36,19 @@ class Gromacs(
# on: turn on, but allow groamcs to disable it if GPU-aware MPI is not supported
# force: turn on and force gromacs to use GPU-aware MPI. May result in error if unsupported
variant(
"gpu-aware-mpi",
"direct-gpu-comm",
default="on",
values=("on", "off", "force"),
description="Use GPU-aware MPI",
)

variant(
"sycl",
default=True,
values=(True, False),
description="Enable GPU-aware MPI",
)

def compute_applications_section(self):
# MPI-only defaults
self.add_experiment_variable("n_ranks", 8, True)
Expand Down Expand Up @@ -86,6 +93,7 @@ def compute_applications_section(self):
"pme": "auto",
"bonded": f"{bonded_target}",
"update": f"{target}",
"additional_args": " -pin {pin} -nb {nb} -pme {pme} -bonded {bonded} -update {update} -maxh {maxh} -nstlist {nstlist} -npme {npme} ",
}

for k, v in input_variables.items():
Expand All @@ -100,14 +108,16 @@ def compute_applications_section(self):
)

def compute_package_section(self):
spack_specs = "~hwloc"
spack_specs += "+sycl" if self.spec.satisfies("+rocm") else "~sycl"
spack_specs = "+mpi~hwloc"
spack_specs += "+sycl" if self.spec.satisfies("+sycl") else "~sycl"

if self.spec.satisfies("+cuda") or self.spec.satisfies("+rocm"):
spack_specs += f" gpu-aware-mpi={self.spec.variants['gpu-aware-mpi'][0]} "
spack_specs += (
f" direct-gpu-comm={self.spec.variants['direct-gpu-comm'][0]} "
)
spack_specs += " ~double "
else:
spack_specs += " gpu-aware-mpi=off "
spack_specs += " direct-gpu-comm=off "

self.add_package_spec(
self.name,
Expand Down
5 changes: 5 additions & 0 deletions modifiers/allocation/modifier.py
Original file line number Diff line number Diff line change
Expand Up @@ -353,6 +353,7 @@ def slurm_instructions(self, v):
sbatch_directives = list(f"#SBATCH {x}" for x in (sbatch_opts))

v.mpi_command = f"{launch_cmd} {' '.join(srun_opts)}"
v.single_rank_mpi_command = f"srun -n 1 -N 1 {'--gpus 1' if v.n_gpus else ''}"

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is single_rank_mpi_command needed here?

v.batch_submit = "sbatch {execute_experiment}"
v.allocation_directives = "\n".join(sbatch_directives)

Expand Down Expand Up @@ -439,6 +440,7 @@ def flux_instructions(self, v):
batch_directives = list(f"# flux: {x}" for x in (batch_opts))

v.mpi_command = f"{launch_cmd} {' '.join([cmd_ranks] + cmd_opts)}"
v.single_rank_mpi_command = f"flux run -n 1 -N 1 {'-g=1' if v.n_gpus else ''}"
v.batch_submit = "flux batch {execute_experiment}"
v.allocation_directives = "\n".join(batch_directives)

Expand All @@ -447,6 +449,7 @@ def mpi_instructions(self, v):
cmd_opts.extend([f"-n {v.n_ranks}"])

v.mpi_command = "mpirun " + " ".join(cmd_opts)
v.single_rank_mpi_command = "mpirun -n 1"
v.batch_submit = "{execute_experiment}"
v.allocation_directives = ""

Expand Down Expand Up @@ -477,6 +480,7 @@ def lsf_instructions(self, v):
batch_directives = list(f"#BSUB {x}" for x in batch_opts)

v.mpi_command = f"lrun {' '.join(cmd_opts)}"
v.single_rank_mpi_command = f"lrun -n 1 -N 1 {'-g 1' if v.n_gpus else ''}"
v.batch_submit = "bsub {execute_experiment}"
v.allocation_directives = "\n".join(batch_directives)

Expand All @@ -493,6 +497,7 @@ def pjm_instructions(self, v):
batch_directives = list(f"#PJM {x}" for x in batch_opts)

v.mpi_command = "mpiexec " + " ".join(cmd_opts)
v.single_rank_mpi_command = "mpiexec --mpi proc=1"
v.batch_submit = "pjsub {execute_experiment}"
v.allocation_directives = "\n".join(batch_directives)

Expand Down
Loading
Loading