Address review comments

Author: BerengerBerthoul

README.md

@@ -1,4 +1,3 @@
-import numpy as np
 import pytest
 import sys
 from mpi4py import MPI
@@ -321,7 +320,7 @@ def schedule_test(item, available_procs, inter_comm):
 
     # mark the procs as busy
     for sub_rank in sub_ranks:
-        available_procs[sub_rank] = False
+        available_procs[sub_rank] = 0
 
     # TODO isend would be slightly better (less waiting)
     for sub_rank in sub_ranks:
@@ -357,15 +356,15 @@ def wait_test_to_complete(items_to_run, session, available_procs, inter_comm):
 
     # the procs are now available
     for sub_rank in sub_ranks:
-        available_procs[sub_rank] = True
+        available_procs[sub_rank] = 1
 
     # "run" the test (i.e. trigger PyTest pipeline but do not really run the code)
     nextitem = None  # not known at this point
     run_item_test(item, nextitem, session)
 
 
 def wait_last_tests_to_complete(items_to_run, session, available_procs, inter_comm):
-    while np.sum(available_procs) < len(available_procs):
+    while sum(available_procs) < len(available_procs):
         wait_test_to_complete(items_to_run, session, available_procs, inter_comm)
 
 
@@ -450,10 +449,10 @@ def pytest_runtestloop(self, session) -> bool:
 
             # schedule tests to run
             items_left_to_run = sorted(items_to_run, key=lambda item: item.n_proc)
-            available_procs = np.ones(n_workers, dtype=np.int8)
+            available_procs = [1] * n_workers
 
             while len(items_left_to_run) > 0:
-                n_av_procs = np.sum(available_procs)
+                n_av_procs = sum(available_procs)
 
                 item_idx = item_with_biggest_admissible_n_proc(items_left_to_run, n_av_procs)
 

pytest_parallel/mpi_reporter.py

@@ -10,6 +10,9 @@
 import pytest
 from _pytest.terminal import TerminalReporter
 
+class PytestParallelError(ValueError):
+  pass
+
 # --------------------------------------------------------------------------
 def pytest_addoption(parser):
     parser.addoption(
@@ -26,9 +29,8 @@ def pytest_addoption(parser):
 
     parser.addoption('--slurm-options', dest='slurm_options', type=str, help='list of SLURM options e.g. "--time=00:30:00 --qos=my_queue"')
     parser.addoption('--slurm-srun-options', dest='slurm_srun_options', type=str, help='list of SLURM srun options e.g. "--mem-per-cpu=4GB"')
-    parser.addoption('--slurm-additional-cmds', dest='slurm_additional_cmds', type=str, help='list of commands to pass to SLURM job e.g. "source my_env.sh"')
+    parser.addoption('--slurm-init-cmds', dest='slurm_init_cmds', type=str, help='list of commands to pass to SLURM job e.g. "source my_env.sh"')
     parser.addoption('--slurm-file', dest='slurm_file', type=str, help='Path to file containing header of SLURM job') # TODO DEL
-    parser.addoption('--slurm-sub-command', dest='slurm_sub_command', type=str, help='SLURM submission command (defaults to `sbatch`)') # TODO DEL
 
     if sys.version_info >= (3,9):
         parser.addoption('--slurm-export-env', dest='slurm_export_env', action=argparse.BooleanOptionalAction, default=True)
@@ -51,19 +53,21 @@ def pytest_addoption(parser):
     #    because it can mess SLURM `srun`
     if "--scheduler=slurm" in sys.argv:
         assert 'mpi4py.MPI' not in sys.modules, 'Internal pytest_parallel error: mpi4py.MPI should not be imported' \
-                                                ' when we are about to register and environment for SLURM' \
+                                                ' when we are about to register an environment for SLURM' \
                                                 ' (because importing mpi4py.MPI makes the current process look like and MPI process,' \
                                                 ' and SLURM does not like that)'
-        assert os.getenv('I_MPI_MPIRUN') is None,  'Internal pytest_parallel error: the environment variable I_MPI_MPIRUN is set' \
-                                                  f' with value "{os.getenv("I_MPI_MPIRUN")}"' \
-                                                   ' while pytest was invoked with "--scheduler=slurm".\n' \
-                                                   ' This indicates that pytest was run through MPI, and SLURM generally does not like that.\n' \
-                                                   ' With "--scheduler=slurm", just run `pytest` directly, not through `mpirun/mpiexec/srun`,\n' \
-                                                   ' because it will launch MPI itself (you may want to use --n-workers=<number of processes>).'
+        if os.getenv('I_MPI_MPIRUN') is not None:
+            err_msg =  'Internal pytest_parallel error: the environment variable I_MPI_MPIRUN is set' \
+                      f' (it has value "{os.getenv("I_MPI_MPIRUN")}"),\n' \
+                       ' while pytest was invoked with "--scheduler=slurm".\n' \
+                       ' This indicates that pytest was run through MPI, and SLURM generally does not like that.\n' \
+                       ' With "--scheduler=slurm", just run `pytest` directly, not through `mpirun/mpiexec/srun`,\n' \
+                       ' because it will launch MPI itself (you may want to use --n-workers=<number of processes>).'
+            raise PytestParallelError(err_msg)
 
         r = subprocess.run(['env','--null'], stdout=subprocess.PIPE) # `--null`: end each output line with NUL, required by `sbatch --export-file`
 
-        assert r.returncode==0, 'SLURM scheduler: error when writing `env` to `pytest_slurm/env_vars.sh`'
+        assert r.returncode==0, 'Internal pytest_parallel SLURM schedule error: error when writing `env` to `pytest_slurm/env_vars.sh`'
         pytest._pytest_parallel_env_vars = r.stdout
 
 # --------------------------------------------------------------------------
@@ -95,40 +99,50 @@ def pytest_configure(config):
     n_workers = config.getoption('n_workers')
     slurm_options = config.getoption('slurm_options')
     slurm_srun_options = config.getoption('slurm_srun_options')
-    slurm_additional_cmds = config.getoption('slurm_additional_cmds')
+    slurm_init_cmds = config.getoption('slurm_init_cmds')
     is_worker = config.getoption('_worker')
     slurm_file = config.getoption('slurm_file')
     slurm_export_env = config.getoption('slurm_export_env')
-    slurm_sub_command = config.getoption('slurm_sub_command')
     detach = config.getoption('detach')
     if scheduler != 'slurm' and scheduler != 'shell':
-        assert not is_worker, 'Option `--slurm-worker` only available when `--scheduler=slurm` or `--scheduler=shell`'
+        assert not is_worker, f'Internal pytest_parallel error `--_worker` not available with`--scheduler={scheduler}`'
     if (scheduler == 'slurm' or scheduler == 'shell') and not is_worker:
-        assert n_workers, f'You need to specify `--n-workers` when `--scheduler={scheduler}`'
+        if n_workers is None:
+            raise PytestParallelError(f'You need to specify `--n-workers` when `--scheduler={scheduler}`')
     if scheduler != 'slurm':
-        assert not slurm_options, 'Option `--slurm-options` only available when `--scheduler=slurm`'
-        assert not slurm_srun_options, 'Option `--slurms-run-options` only available when `--scheduler=slurm`'
-        assert not slurm_additional_cmds, 'Option `--slurm-additional-cmds` only available when `--scheduler=slurm`'
-        assert not slurm_file, 'Option `--slurm-file` only available when `--scheduler=slurm`'
+        if slurm_options is not None:
+            raise PytestParallelError('Option `--slurm-options` only available when `--scheduler=slurm`')
+        if slurm_srun_options is not None:
+            raise PytestParallelError('Option `--slurms-run-options` only available when `--scheduler=slurm`')
+        if slurm_init_cmds is not None:
+            raise PytestParallelError('Option `--slurm-init-cmds` only available when `--scheduler=slurm`')
+        if slurm_file is not None:
+            raise PytestParallelError('Option `--slurm-file` only available when `--scheduler=slurm`')
 
     if (scheduler == 'shell' or scheduler == 'slurm') and not is_worker:
         from mpi4py import MPI
-        assert MPI.COMM_WORLD.size == 1, 'Do not launch `pytest_parallel` on more that one process\n' \
-                                         'when `--scheduler=shell` or `--scheduler=slurm`.\n' \
-                                         '`pytest_parallel` spawns MPI processes itself.\n' \
-                                         f'You may want to use --n-workers={MPI.COMM_WORLD.size}.'
+        if MPI.COMM_WORLD.size != 1:
+            err_msg = 'Do not launch `pytest_parallel` on more that one process when `--scheduler=shell` or `--scheduler=slurm`.\n' \
+                      '`pytest_parallel` will spawn MPI processes itself.\n' \
+                     f'You may want to use --n-workers={MPI.COMM_WORLD.size}.'
+            raise PytestParallelError(err_msg)
 
 
 
     if scheduler == 'slurm' and not is_worker:
-        assert slurm_options or slurm_file, 'You need to specify either `--slurm-options` or `--slurm-file` when `--scheduler=slurm`'
+        if slurm_options is None and slurm_file is None:
+            raise PytestParallelError('You need to specify either `--slurm-options` or `--slurm-file` when `--scheduler=slurm`')
         if slurm_options:
-            assert not slurm_file, 'You need to specify either `--slurm-options` or `--slurm-file`, but not both'
+            if slurm_file:
+                raise PytestParallelError('You need to specify either `--slurm-options` or `--slurm-file`, but not both')
         if slurm_file:
-            assert not slurm_options, 'You need to specify either `--slurm-options` or `--slurm-file`, but not both'
-            assert not slurm_additional_cmds, 'You cannot specify `--slurm-additional-cmds` together with `--slurm-file`'
+            if slurm_options:
+                raise PytestParallelError('You need to specify either `--slurm-options` or `--slurm-file`, but not both')
+            if slurm_init_cmds:
+                raise PytestParallelError('You cannot specify `--slurm-init-cmds` together with `--slurm-file`')
 
-        assert '-n=' not in slurm_options and '--ntasks=' not in slurm_options, 'Do not specify `-n/--ntasks` in `--slurm-options` (it is deduced from the `--n-worker` value).'
+        if '-n=' in slurm_options or '--ntasks=' in slurm_options:
+                raise PytestParallelError('Do not specify `-n/--ntasks` in `--slurm-options` (it is deduced from the `--n-worker` value).')
 
         from .slurm_scheduler import SlurmScheduler
 
@@ -143,12 +157,11 @@ def pytest_configure(config):
           main_invoke_params = main_invoke_params.replace(file_or_dir, '')
         slurm_option_list = slurm_options.split() if slurm_options is not None else []
         slurm_conf = {
-            'options'        : slurm_option_list,
-            'srun_options'   : slurm_srun_options,
-            'additional_cmds': slurm_additional_cmds,
-            'file'           : slurm_file,
-            'export_env'     : slurm_export_env,
-            'sub_command'    : slurm_sub_command,
+            'options'     : slurm_option_list,
+            'srun_options': slurm_srun_options,
+            'init_cmds'   : slurm_init_cmds,
+            'file'        : slurm_file,
+            'export_env'  : slurm_export_env,
         }
         plugin = SlurmScheduler(main_invoke_params, n_workers, slurm_conf, detach)
 

pytest_parallel/plugin.py

@@ -26,25 +26,26 @@ def _file_path(when):
 test_info = {'test_idx': args._test_idx, 'fatal_error': None} # TODO no fatal_error=None (absense means no error)
 
 # 'fatal_error' file
-try:
-  file_path = _file_path('fatal_error')
+file_path = _file_path('fatal_error')
+if file_path.exists():
   with open(file_path, 'r') as file:
     fatal_error = file.read()
     test_info['fatal_error'] = fatal_error
-except FileNotFoundError: # There was no fatal error
-  pass
 
 
 # 'setup/call/teardown' files
 already_failed = False
 for when in ('setup', 'call', 'teardown'):
-  try:
-    file_path = _file_path(when)
-    with open(file_path, 'rb') as file:
-      report_info = file.read()
-      report_info = pickle.loads(report_info)
-      test_info[when] = report_info
-  except FileNotFoundError: # Supposedly not found because the test crashed before writing the file
+  file_path = _file_path(when)
+  if file_path.exists():
+    try:
+      with open(file_path, 'rb') as file:
+        report_info = file.read()
+        report_info = pickle.loads(report_info)
+        test_info[when] = report_info
+    except pickle.PickleError:
+      test_info['fatal_error'] = f'FATAL ERROR in pytest_parallel : unable to decode {file_path}'
+  else: # Supposedly not found because the test crashed before writing the file
     collect_longrepr = []
     msg = f'Error: the test crashed. '
     red = 31
@@ -63,8 +64,6 @@ def _file_path(when):
                        'duration': 0, } # unable to report accurately
 
     already_failed = True
-  except pickle.PickleError:
-    test_info['fatal_error'] = f'FATAL ERROR in pytest_parallel : unable to decode {file_path}'
 
 
 with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:

pytest_parallel/send_report.py

@@ -5,13 +5,13 @@
 import socket
 import pickle
 from pathlib import Path
-from .utils.socket import recv as socket_recv, setup_socket
+from .utils.socket import recv as socket_recv
+from .utils.socket import setup_socket
 from .utils.items import get_n_proc_for_test, add_n_procs, run_item_test, mark_original_index, mark_skip
 from .utils.file import remove_exotic_chars, create_folders
 from .algo import partition
 from .static_scheduler_utils import group_items_by_parallel_steps
 from mpi4py import MPI
-import numpy as np
 
 def mpi_command(current_proc, n_proc):
     mpi_vendor = MPI.get_vendor()[0]
@@ -43,7 +43,7 @@ def submit_items(items_to_run, SCHEDULER_IP_ADDRESS, port, session_folder, main_
         cmd += mpi_command(current_proc, item.n_proc)
         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' ; python -m pytest_parallel.send_report {socket_flags} --_test_idx={test_idx} --_test_name={test_out_file}'
+        cmd += f' ; python3 -m pytest_parallel.send_report {socket_flags} --_test_idx={test_idx} --_test_name={test_out_file}'
         cmd += ')'
         cmds.append(cmd)
         current_proc += item.n_proc
@@ -74,18 +74,17 @@ def submit_items(items_to_run, SCHEDULER_IP_ADDRESS, port, session_folder, main_
 
 def receive_items(items, session, socket, n_item_to_recv):
     # > Precondition: Items must keep their original order to pick up the right item at the reception
-    original_indices = np.array([item.original_index for item in items])
-    assert (original_indices==np.arange(len(items))).all()
+    original_indices = [item.original_index for item in items]
+    assert original_indices==list(range(len(items)))
 
     while n_item_to_recv>0:
         conn, addr = socket.accept()
         with conn:
             msg = socket_recv(conn)
         test_info = pickle.loads(msg) # the worker is supposed to have send a dict with the correct structured information
-        #print(f"{test_info=}")
         if 'signal_info' in test_info:
             print('signal_info= ',test_info['signal_info'])
-            break;
+            break
         else:
             test_idx = test_info['test_idx']
             if test_info['fatal_error'] is not None:

pytest_parallel/shell_static_scheduler.py

@@ -3,19 +3,12 @@
 import socket
 import pickle
 from pathlib import Path
-from .utils.socket import recv as socket_recv, setup_socket
+from .utils.socket import recv as socket_recv
+from .utils.socket import setup_socket
 from .utils.items import get_n_proc_for_test, add_n_procs, run_item_test, mark_original_index, mark_skip
 from .utils.file import remove_exotic_chars, create_folders
 from .algo import partition
 
-def parse_job_id_from_submission_output(s):
-    # At this point, we are trying to guess -_-
-    # Here we supposed that the command for submitting the job
-    #    returned string with only one number,
-    #    and that this number is the job id
-    import re
-    return int(re.search(r'\d+', str(s)).group())
-
 def submit_items(items_to_run, socket, session_folder, main_invoke_params, ntasks, slurm_conf):
     SCHEDULER_IP_ADDRESS, port = setup_socket(socket)
 
@@ -45,8 +38,8 @@ def submit_items(items_to_run, socket, session_folder, main_invoke_params, ntask
       srun_options = ''
     socket_flags = f"--_scheduler_ip_address={SCHEDULER_IP_ADDRESS} --_scheduler_port={port} --_session_folder={session_folder}"
     cmds = ''
-    if slurm_conf['additional_cmds'] is not None:
-        cmds += slurm_conf['additional_cmds'] + '\n'
+    if slurm_conf['init_cmds'] is not None:
+        cmds += slurm_conf['init_cmds'] + '\n'
     for item in items:
         test_idx = item.original_index
         test_out_file = f'.pytest_parallel/{session_folder}/{remove_exotic_chars(item.nodeid)}'
@@ -58,7 +51,7 @@ def submit_items(items_to_run, socket, session_folder, main_invoke_params, ntask
         cmd +=  ' -l' # 
         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' ; python -m pytest_parallel.send_report {socket_flags} --_test_idx={test_idx} --_test_name={test_out_file}'
+        cmd += f' ; python3 -m pytest_parallel.send_report {socket_flags} --_test_idx={test_idx} --_test_name={test_out_file}'
         cmd +=  ')'
         cmd +=  ' &\n' # launch everything in parallel
         cmds += cmd
@@ -73,23 +66,17 @@ def submit_items(items_to_run, socket, session_folder, main_invoke_params, ntask
     with open(f'.pytest_parallel/{session_folder}/env_vars.sh','wb') as f:
       f.write(pytest._pytest_parallel_env_vars)
 
-    if slurm_conf['sub_command'] is None:
-        if slurm_conf['export_env']:
-            sbatch_cmd = f'sbatch --parsable --export-file=.pytest_parallel/{session_folder}/env_vars.sh .pytest_parallel/{session_folder}/job.sh'
-        else:
-            sbatch_cmd = f'sbatch --parsable .pytest_parallel/{session_folder}/job.sh'
+    if slurm_conf['export_env']:
+        sbatch_cmd = f'sbatch --parsable --export-file=.pytest_parallel/{session_folder}/env_vars.sh .pytest_parallel/{session_folder}/job.sh'
     else:
-        sbatch_cmd = slurm_conf['sub_command'] + ' .pytest_parallel/{session_folder}/job.sh'
+        sbatch_cmd = f'sbatch --parsable .pytest_parallel/{session_folder}/job.sh'
 
     p = subprocess.Popen([sbatch_cmd], shell=True, stdout=subprocess.PIPE)
     print('\nSubmitting tests to SLURM...')
     returncode = p.wait()
     assert returncode==0, f'Error when submitting to SLURM with `{sbatch_cmd}`'
 
-    if slurm_conf['sub_command'] is None:
-        slurm_job_id = int(p.stdout.read())
-    else:
-        slurm_job_id = parse_job_id_from_submission_output(p.stdout.read())
+    slurm_job_id = int(p.stdout.read())
 
     print(f'SLURM job {slurm_job_id} has been submitted')
     return slurm_job_id

pytest_parallel/slurm_scheduler.py

@@ -13,6 +13,6 @@ def remove_exotic_chars(s):
 
 def create_folders():
   Path('.pytest_parallel').mkdir(exist_ok=True)
-  session_folder_abs = Path(tempfile.mkdtemp(dir='.pytest_parallel'))
+  session_folder_abs = Path(tempfile.mkdtemp(dir='.pytest_parallel')) # create a folder that did not already exist
   Path(session_folder_abs/'_partial').mkdir()
   return session_folder_abs.name