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

@@ -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 = (

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]:

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)