Coverage for /builds/ase/ase/ase/calculators/exciting/runner.py: 60.00%

1# fmt: off 

2 

3"""Binary runner and results class.""" 

4import os 

5import subprocess 

6import time 

7from pathlib import Path 

8from typing import List, Optional, Union 

9 

10 

11class SubprocessRunResults: 

12 """Results returned from subprocess.run().""" 

13 

14 def __init__( 

15 self, stdout, stderr, return_code: int, 

16 process_time: Optional[float] = None): 

17 self.stdout = stdout 

18 self.stderr = stderr 

19 self.return_code = return_code 

20 self.success = return_code == 0 

21 self.process_time = process_time 

22 

23 

24class SimpleBinaryRunner: 

25 """Class to execute a subprocess.""" 

26 path_type = Union[str, Path] 

27 

28 def __init__(self, 

29 binary, 

30 run_argv: List[str], 

31 omp_num_threads: int, 

32 directory: path_type = './', 

33 args=None) -> None: 

34 """Initialise class. 

35 

36 :param binary: Binary name prepended by full path, or just binary name 

37 (if present in $PATH). 

38 :param run_argv: Run commands sequentially as a list of str. 

39 For example: 

40 * For serial: ['./'] or [''] 

41 * For MPI: ['mpirun', '-np', '2'] 

42 :param omp_num_threads: Number of OMP threads. 

43 :param args: Optional arguments for the binary. 

44 """ 

45 if args is None: 

46 args = [] 

47 self.binary = binary 

48 self.directory = directory 

49 

50 self.run_argv = run_argv 

51 

52 self.omp_num_threads = omp_num_threads 

53 self.args = args 

54 

55 if directory is not None and not Path(directory).is_dir(): 

56 raise OSError(f"Run directory does not exist: {directory}") 

57 

58 if omp_num_threads <= 0: 

59 raise ValueError("Number of OMP threads must be > 0") 

60 

61 def compose_execution_list(self) -> list: 

62 """Generate a complete list of strings to pass to subprocess.run(). 

63 

64 This is done to execute the calculation. 

65 

66 For example, given: 

67 ['mpirun', '-np, '2'] + ['binary.exe'] + ['>', 'std.out'] 

68 

69 return ['mpirun', '-np, '2', 'binary.exe', '>', 'std.out'] 

70 """ 

71 return self.run_argv + [self.binary] + self.args 

72 

73 def run(self) -> SubprocessRunResults: 

74 """Run a binary.""" 

75 execution_list = self.compose_execution_list() 

76 my_env = {**os.environ} 

77 

78 time_start: float = time.time() 

79 result = subprocess.run(execution_list, 

80 env=my_env, 

81 capture_output=True, 

82 cwd=self.directory, check=False) 

83 total_time = time.time() - time_start 

84 return SubprocessRunResults( 

85 result.stdout, result.stderr, result.returncode, total_time)