Newer
Older
import sys
import os
from pathlib import Path
def define_python_interpreter(python_path=None, conda_env=None):
conda_base = "module load Anaconda3\nconda activate {conda_env}"
venv_base = "source {python_path}"
if conda_env is not None:
env = conda_base.format(conda_env=conda_env)
elif python_path is not None:
parent = Path(python_path).absolute().parent
env = venv_base.format(python_path=parent.joinpath('activate'))
else:
conda_env = os.environ.get('CONDA_PREFIX')
if conda_env:
env = conda_base.format(conda_env=conda_env)
else:
parent = Path(sys.executable).absolute().parent
env = venv_base.format(python_path=parent.joinpath('activate'))
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
return env
class CustomHelpFormatter(argparse.MetavarTypeHelpFormatter):
def add_arguments(self, actions):
# Sort actions by their group title
actions = sorted(actions, key=lambda x: x.container.title if x.container.title else '')
super(CustomHelpFormatter, self).add_arguments(actions)
def batch_parser(
cpus_per_task: int | None = None,
gpus: int | None = None,
partition: str | None = None,
mem: str | None = None,
time: str | None = '12:00:00',
reservation: str | None = None,
slurm_log_dir: str | Path | None = './out',
**kwargs
) -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(add_help=False,
formatter_class=CustomHelpFormatter)
slurm = parser.add_argument_group(title='Slurm Options')
slurm.add_argument('--batch', action='store_true', default=False,
help="Convert as a batch array job.")
slurm.add_argument('-n', '--ntasks', type=int, default=1)
slurm.add_argument('-c', '--cpus-per-task', type=int, default=cpus_per_task)
slurm.add_argument('-g', '--gpus', type=int, default=gpus, choices=[0, 1])
slurm.add_argument('-p', '--partition', type=str, default=partition)
slurm.add_argument('-t', '--time', type=str, default=time)
slurm.add_argument('-m', '--mem', type=str, default=mem)
slurm.add_argument('--reservation',type=str,default=reservation)
slurm.add_argument('--slurm-log-dir', type=Path, default=slurm_log_dir,
help='Output log directory. If the directory does not exist, it will be created automatically.')
return parser
def setup_slurm_logs(slurm_log_dir,log_basename):
slurm_log_dir = slurm_log_dir.absolute()
slurm_log_dir.mkdir(exist_ok = True,parents=True,mode = 0o2770)
out_log,err_log = [str(slurm_log_dir.joinpath(f'{log_basename}_%A_%a.out')),str(slurm_log_dir.joinpath(f'{log_basename}_%A_%a.err'))]
slurm_logs = {'output_log':out_log,'error_log':err_log}
return slurm_logs