Coverage for src/shephex/executor/slurm/slurm_header.py: 100%

41 statements  

« prev     ^ index     » next       coverage.py v7.6.1, created at 2025-03-29 18:45 +0100

1from copy import deepcopy 

2from typing import Tuple 

3 

4from shephex.executor.slurm.slurm_options import valid_options 

5 

6flags = [str(option[0]).replace('-', '') for option in valid_options] 

7long_flags = [option[1].replace('--', '') for option in valid_options] 

8 

9 

10class HeaderOption: 

11 def __init__(self, key: str, value: str) -> None: 

12 """ 

13 A single header option for a Slurm script 

14 Slurm script, e.g. 

15 #SBATCH --partition=<partition> 

16 

17 Checked for correctness against 'valid_options'. 

18 

19 Parameters 

20 ---------- 

21 key : str 

22 Option key, such as 'partition' 

23 value : str 

24 Option value, such as 'gpu' 

25 """ 

26 

27 self.key, self.option_index = self.get_name_and_index(key) 

28 self.value = value 

29 

30 def __repr__(self) -> str: 

31 return f'#SBATCH --{self.key}={self.value}' 

32 

33 def __str__(self) -> str: 

34 return self.__repr__() 

35 

36 @staticmethod 

37 def get_name_and_index(item: str) -> Tuple[str, int]: 

38 item = item.replace('_', '-') 

39 if item in flags: 

40 index = flags.index(item) 

41 return long_flags[index], index 

42 elif item in long_flags: 

43 index = long_flags.index(item) 

44 return item, index 

45 else: 

46 raise ValueError(f'Invalid option: {item}') 

47 

48 

49class SlurmHeader: 

50 def __init__(self) -> None: 

51 """ 

52 Header portion of a Slurm script, consisting of a list of header options. 

53 """ 

54 self.options = [] 

55 

56 def add(self, key: str, value: str) -> None: 

57 """ 

58 Add a header option to the header. 

59 Raises an error if the option already exists. 

60 

61 Parameters 

62 ---------- 

63 key : str 

64 Option key, such as 'partition' 

65 """ 

66 option = HeaderOption(key, value) 

67 

68 for i, existing_option in enumerate(self.options): 

69 if option.option_index == existing_option.option_index: 

70 raise ValueError( 

71 f"Option '{key}' already exists in header options ({existing_option.key})" 

72 ) 

73 

74 self.options.append(option) 

75 

76 def __repr__(self) -> str: 

77 s = '#!/bin/sh' 

78 for option in self.options: 

79 s += '\n' + str(option) 

80 return s 

81 

82 def copy(self) -> 'SlurmHeader': 

83 return deepcopy(self) 

84 

85 def to_dict(self) -> dict: 

86 return {option.key: option.value for option in self.options}