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
« prev ^ index » next coverage.py v7.6.1, created at 2025-03-29 18:45 +0100
1from copy import deepcopy
2from typing import Tuple
4from shephex.executor.slurm.slurm_options import valid_options
6flags = [str(option[0]).replace('-', '') for option in valid_options]
7long_flags = [option[1].replace('--', '') for option in valid_options]
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>
17 Checked for correctness against 'valid_options'.
19 Parameters
20 ----------
21 key : str
22 Option key, such as 'partition'
23 value : str
24 Option value, such as 'gpu'
25 """
27 self.key, self.option_index = self.get_name_and_index(key)
28 self.value = value
30 def __repr__(self) -> str:
31 return f'#SBATCH --{self.key}={self.value}'
33 def __str__(self) -> str:
34 return self.__repr__()
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}')
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 = []
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.
61 Parameters
62 ----------
63 key : str
64 Option key, such as 'partition'
65 """
66 option = HeaderOption(key, value)
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 )
74 self.options.append(option)
76 def __repr__(self) -> str:
77 s = '#!/bin/sh'
78 for option in self.options:
79 s += '\n' + str(option)
80 return s
82 def copy(self) -> 'SlurmHeader':
83 return deepcopy(self)
85 def to_dict(self) -> dict:
86 return {option.key: option.value for option in self.options}