From c6ec746be1d6e2a39deba10500df5bf828fc9fc3 Mon Sep 17 00:00:00 2001 From: Riccardo Balin Date: Wed, 18 Oct 2023 13:28:29 -0400 Subject: [PATCH] Adding affinity scripts to PALS mpiexec settings (#282) Added an option to specify an affinity script for the PALSMpiexecSettings and add it to the run command after the run arguments, with no "--" prefix. Optional arguments can be passed after the affinity script. [ committed by @rickybalin ] [ reviewed by @al-rigazzi ] --- smartsim/settings/mpiSettings.py | 1 + smartsim/settings/palsSettings.py | 14 ++++++++++++++ tests/test_pals_settings.py | 5 +++++ 3 files changed, 20 insertions(+) diff --git a/smartsim/settings/mpiSettings.py b/smartsim/settings/mpiSettings.py index 4a3a5afc4..5b6b520e3 100644 --- a/smartsim/settings/mpiSettings.py +++ b/smartsim/settings/mpiSettings.py @@ -81,6 +81,7 @@ def __init__( **kwargs, ) self.mpmd: t.List[RunSettings] = [] + self.affinity_script: t.List[str] = [] if not shutil.which(self._run_command): msg = ( diff --git a/smartsim/settings/palsSettings.py b/smartsim/settings/palsSettings.py index 3a8868793..1b1055976 100644 --- a/smartsim/settings/palsSettings.py +++ b/smartsim/settings/palsSettings.py @@ -182,6 +182,16 @@ def set_walltime(self, walltime: str) -> None: """ logger.warning("set_walltime not supported under PALS") + def set_gpu_affinity_script(self, affinity: str, *args: t.Any) -> None: + """ Set the GPU affinity through a bash script + + :param affinity: path to the affinity script + :type affinity: str + """ + self.affinity_script.append(str(affinity)) + for arg in args: + self.affinity_script.append(str(arg)) + def format_run_args(self) -> t.List[str]: """Return a list of MPI-standard formatted run arguments @@ -199,6 +209,10 @@ def format_run_args(self) -> t.List[str]: args += [prefix + opt] else: args += [prefix + opt, str(value)] + + if self.affinity_script: + args += self.affinity_script + return args def format_env_vars(self) -> t.List[str]: diff --git a/tests/test_pals_settings.py b/tests/test_pals_settings.py index 7bc3a6520..4c837149d 100644 --- a/tests/test_pals_settings.py +++ b/tests/test_pals_settings.py @@ -54,6 +54,11 @@ # with pytest.raises(SSUnsupportedError): # func(None) +def test_affinity_script(): + settings = PalsMpiexecSettings(default_exe, **default_kwargs) + settings.set_gpu_affinity_script("/path/to/set_affinity_gpu.sh", 1, 2) + assert settings.format_run_args() == ["/path/to/set_affinity_gpu.sh", "1", "2"] + def test_cpu_binding_type(): settings = PalsMpiexecSettings(default_exe, **default_kwargs)