Coverage for gpaw/cli/sbatch.py: 30%

1import subprocess 

2import sys 

3import os 

4 

5usage = """gpaw sbatch [-0] -- [sbatch options] script.py [script options] 

6 or: gpaw sbatch [-0] -- [sbatch options] python -m module [module options] 

7""" 

8 

9 

10class CLICommand: 

11 """Submit a GPAW Python script via sbatch. 

12 

13 

14 If a virtual environment is active when submitting, it will be activated for 

15 the job as well. 

16 """ 

17 

18 @staticmethod 

19 def add_arguments(parser): 

20 parser.usage = usage 

21 parser.add_argument('-0', '--test', action='store_true', 

22 help='Dry run: Print driver script.') 

23 parser.add_argument('arguments', nargs='*') 

24 

25 @staticmethod 

26 def run(args): 

27 script = '#!/bin/bash -l\n' 

28 for i, arg in enumerate(args.arguments): 

29 if arg.endswith('.py'): 

30 break 

31 else: 

32 for i, arg in enumerate(args.arguments): 

33 if (arg.startswith('python') and 

34 len(args.arguments) > i + 1 and 

35 args.arguments[i + 1].startswith('-m')): 

36 del args.arguments[i] 

37 break 

38 else: 

39 print('No script.py found!', file=sys.stderr) 

40 return 

41 

42 if arg.endswith('.py'): 

43 for line in open(arg): 

44 if line.startswith('#SBATCH'): 

45 script += line 

46 script += ('cd $SLURM_SUBMIT_DIR\n') 

47 venv = os.getenv('VIRTUAL_ENV') 

48 if venv: 

49 print('Detected virtual environment:', venv) 

50 script += f'source {venv}/bin/activate\n' 

51 script += ('OMP_NUM_THREADS=1 ' 

52 'mpiexec `echo $GPAW_MPI_OPTIONS` gpaw python {}\n' 

53 .format(' '.join(args.arguments[i:]))) 

54 cmd = ['sbatch'] + args.arguments[:i] 

55 if args.test: 

56 print('sbatch command:') 

57 print(' '.join(cmd)) 

58 print('\nscript:') 

59 print(script, end='') 

60 else: 

61 p = subprocess.Popen(cmd, stdin=subprocess.PIPE, env=os.environ) 

62 p.communicate(script.encode())