From 3036a8d4c0c0f87b1983bebd148c2245c8be2ca0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?B=C3=A9renger=20Berthoul?= Date: Thu, 23 May 2024 14:27:38 +0200 Subject: [PATCH] new option --[no]-slurm-export-env. Useful if --export-file is not supported --- pytest_parallel/plugin.py | 14 ++++++++++++-- pytest_parallel/process_scheduler.py | 5 ++++- 2 files changed, 16 insertions(+), 3 deletions(-) diff --git a/pytest_parallel/plugin.py b/pytest_parallel/plugin.py index 4b78842..a94f2e6 100644 --- a/pytest_parallel/plugin.py +++ b/pytest_parallel/plugin.py @@ -6,6 +6,7 @@ import tempfile import pytest from pathlib import Path +import argparse # -------------------------------------------------------------------------- @@ -23,8 +24,15 @@ 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 --n_tasks=4"') 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-file', dest='slurm_file', type=str, help='Path to file containing header of SLURM job') - parser.addoption('--slurm-sub-command', dest='slurm_sub_command', type=str, help='SLURM submission command (defaults to `sbatch`)') + 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) + else: + parser.addoption('--slurm-export-env', dest='slurm_export_env', action='store_true') + parser.addoption('--no-slurm-export-env', dest='slurm_export_env', action='store_false') + parser.set_defaults(slurm_export_env=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)') @@ -59,6 +67,7 @@ def pytest_configure(config): slurm_additional_cmds = config.getoption('slurm_additional_cmds') slurm_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': @@ -97,6 +106,7 @@ def pytest_configure(config): 'options' : slurm_option_list, 'additional_cmds': slurm_additional_cmds, 'file' : slurm_file, + 'export_env' : slurm_export_env, 'sub_command' : slurm_sub_command, } plugin = ProcessScheduler(main_invoke_params, n_workers, slurm_conf, detach) diff --git a/pytest_parallel/process_scheduler.py b/pytest_parallel/process_scheduler.py index e84aa38..edfe0ee 100644 --- a/pytest_parallel/process_scheduler.py +++ b/pytest_parallel/process_scheduler.py @@ -114,7 +114,10 @@ def submit_items(items_to_run, socket, main_invoke_params, slurm_ntasks, slurm_c f.write(pytest._pytest_parallel_env_vars) if slurm_conf['sub_command'] is None: - sbatch_cmd = 'sbatch --parsable --export-file=pytest_slurm/env_vars.sh pytest_slurm/job.sh' + if slurm_conf['export_env']: + sbatch_cmd = 'sbatch --parsable --export-file=pytest_slurm/env_vars.sh pytest_slurm/job.sh' + else: + sbatch_cmd = 'sbatch --parsable pytest_slurm/job.sh' else: sbatch_cmd = slurm_conf['sub_command'] + ' pytest_slurm/job.sh'