shell scheduler: --use-srun option

Author: BerengerBerthoul

pytest_parallel/plugin.py

@@ -35,9 +35,11 @@ def pytest_addoption(parser):
 
     if sys.version_info >= (3,9):
         parser.addoption('--slurm-export-env', dest='slurm_export_env', action=argparse.BooleanOptionalAction, default=True)
+        parser.addoption('--use-srun', dest='use_srun', action=argparse.BooleanOptionalAction, default=None, help='Launch MPI processes through srun (only possible when `--scheduler=shell`')
     else:
         parser.addoption('--slurm-export-env', dest='slurm_export_env', default=False, action='store_true')
         parser.addoption('--no-slurm-export-env', dest='slurm_export_env', action='store_false')
+        parser.addoption('--use-srun', dest='use_srun', default=None, action='store_true')
 
     parser.addoption('--detach', dest='detach', action='store_true', help='Detach SLURM jobs: do not send reports to the scheduling process (useful to launch slurm job.sh separately)')
 
@@ -104,7 +106,11 @@ def pytest_configure(config):
     is_worker = config.getoption('_worker')
     slurm_file = config.getoption('slurm_file')
     slurm_export_env = config.getoption('slurm_export_env')
+    use_srun = config.getoption('use_srun')
     detach = config.getoption('detach')
+    if scheduler != 'shell':
+        if use_srun is not None:
+            raise PytestParallelError('Option `--use-srun` only available when `--scheduler=shell`')
     if not scheduler in ['slurm', 'shell']:
         assert not is_worker, f'Internal pytest_parallel error `--_worker` not available with`--scheduler={scheduler}`'
         assert not n_workers, f'pytest_parallel error `--n-workers` not available with`--scheduler={scheduler}`. Launch with `mpirun -np {n_workers}` to run in parallel'
@@ -175,7 +181,7 @@ def pytest_configure(config):
         main_invoke_params = _invoke_params(config.invocation_params.args)
         for file_or_dir in config.option.file_or_dir:
             main_invoke_params = main_invoke_params.replace(file_or_dir, '')
-        plugin = ShellStaticScheduler(main_invoke_params, n_workers, detach)
+        plugin = ShellStaticScheduler(main_invoke_params, n_workers, detach, use_srun)
     else:
         from mpi4py import MPI
         from .mpi_reporter import SequentialScheduler, StaticScheduler, DynamicScheduler

pytest_parallel/shell_static_scheduler.py

@@ -13,19 +13,22 @@
 from .utils.file import remove_exotic_chars, create_folders
 from .static_scheduler_utils import group_items_by_parallel_steps
 
-def mpi_command(current_proc, n_proc):
-    mpi_vendor = MPI.get_vendor()[0]
-    if mpi_vendor == 'Intel MPI':
-        cmd =  f'I_MPI_PIN_PROCESSOR_LIST={current_proc}-{current_proc+n_proc-1}; '
-        cmd += f'mpiexec -np {n_proc}'
-        return cmd
-    elif mpi_vendor == 'Open MPI':
-        cores = ','.join([str(i) for i in range(current_proc,current_proc+n_proc)])
-        return f'mpiexec --cpu-list {cores} -np {n_proc}'
+def mpi_command(current_proc, n_proc, use_srun):
+    if use_srun:
+        return f'srun --exact --ntasks={n_proc}'
     else:
-        assert 0, f'Unknown MPI implementation "{mpi_vendor}"'
+        mpi_vendor = MPI.get_vendor()[0]
+        if mpi_vendor == 'Intel MPI':
+            cmd =  f'I_MPI_PIN_PROCESSOR_LIST={current_proc}-{current_proc+n_proc-1}; '
+            cmd += f'mpiexec -np {n_proc}'
+            return cmd
+        elif mpi_vendor == 'Open MPI':
+            cores = ','.join([str(i) for i in range(current_proc,current_proc+n_proc)])
+            return f'mpiexec --cpu-list {cores} -np {n_proc}'
+        else:
+            assert 0, f'Unknown MPI implementation "{mpi_vendor}"'
 
-def submit_items(items_to_run, SCHEDULER_IP_ADDRESS, port, session_folder, main_invoke_params, i_step, n_step):
+def submit_items(items_to_run, SCHEDULER_IP_ADDRESS, port, session_folder, main_invoke_params, use_srun, i_step, n_step):
     # sort item by comm size to launch bigger first (Note: in case SLURM prioritize first-received items)
     items = sorted(items_to_run, key=lambda item: item.n_proc, reverse=True)
 
@@ -40,7 +43,7 @@ def submit_items(items_to_run, SCHEDULER_IP_ADDRESS, port, session_folder, main_
         test_idx = item.original_index
         test_out_file = f'.pytest_parallel/{session_folder}/{remove_exotic_chars(item.nodeid)}'
         cmd = '('
-        cmd += mpi_command(current_proc, item.n_proc)
+        cmd += mpi_command(current_proc, item.n_proc, use_srun)
         cmd += f' python3 -u -m pytest -s --_worker {socket_flags} {main_invoke_params} --_test_idx={test_idx} {item.config.rootpath}/{item.nodeid}'
         cmd += f' > {test_out_file} 2>&1'
         cmd += f' ; python3 -m pytest_parallel.send_report {socket_flags} --_test_idx={test_idx} --_test_name={test_out_file}'
@@ -99,10 +102,11 @@ def receive_items(items, session, socket, n_item_to_recv):
             n_item_to_recv -= 1
 
 class ShellStaticScheduler:
-    def __init__(self, main_invoke_params, ntasks, detach):
+    def __init__(self, main_invoke_params, ntasks, detach, use_srun):
         self.main_invoke_params = main_invoke_params
         self.ntasks             = ntasks
         self.detach             = detach
+        self.use_srun           = use_srun
 
         self.socket             = socket.socket(socket.AF_INET, socket.SOCK_STREAM) # TODO close at the end
 
@@ -148,7 +152,7 @@ def pytest_runtestloop(self, session) -> bool:
         n_step = len(items_by_steps)
         for i_step,items in enumerate(items_by_steps):
             n_item_to_receive = len(items)
-            sub_process = submit_items(items, SCHEDULER_IP_ADDRESS, port, session_folder, self.main_invoke_params, i_step, n_step)
+            sub_process = submit_items(items, SCHEDULER_IP_ADDRESS, port, session_folder, self.main_invoke_params, self.use_srun, i_step, n_step)
             if not self.detach: # The job steps are supposed to send their reports
                 receive_items(session.items, session, self.socket, n_item_to_receive)
             returncode = sub_process.wait() # at this point, the sub-process should be done since items have been received