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

1"""Binary runner and results class.""" 

2import os 

3import subprocess 

4import time 

5from pathlib import Path 

6from typing import List, Optional, Union 

7 

8 

9class SubprocessRunResults: 

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

11 

12 def __init__( 

13 self, stdout, stderr, return_code: int, 

14 process_time: Optional[float] = None): 

15 self.stdout = stdout 

16 self.stderr = stderr 

17 self.return_code = return_code 

18 self.success = return_code == 0 

19 self.process_time = process_time 

20 

21 

22class SimpleBinaryRunner: 

23 """Class to execute a subprocess.""" 

24 path_type = Union[str, Path] 

25 

26 def __init__(self, 

27 binary, 

28 run_argv: List[str], 

29 omp_num_threads: int, 

30 directory: path_type = './', 

31 args=None) -> None: 

32 """Initialise class. 

33 

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

35 (if present in $PATH). 

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

37 For example: 

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

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

40 :param omp_num_threads: Number of OMP threads. 

41 :param args: Optional arguments for the binary. 

42 """ 

43 if args is None: 

44 args = [] 

45 self.binary = binary 

46 self.directory = directory 

47 

48 self.run_argv = run_argv 

49 

50 self.omp_num_threads = omp_num_threads 

51 self.args = args 

52 

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

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

55 

56 if omp_num_threads <= 0: 

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

58 

59 def compose_execution_list(self) -> list: 

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

61 

62 This is done to execute the calculation. 

63 

64 For example, given: 

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

66 

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

68 """ 

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

70 

71 def run(self) -> SubprocessRunResults: 

72 """Run a binary.""" 

73 execution_list = self.compose_execution_list() 

74 my_env = {**os.environ} 

75 

76 time_start: float = time.time() 

77 result = subprocess.run(execution_list, 

78 env=my_env, 

79 capture_output=True, 

80 cwd=self.directory) 

81 total_time = time.time() - time_start 

82 return SubprocessRunResults( 

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