Coverage for parallel_bilby/parser/generation.py: 96%

1import bilby 

2import bilby_pipe.data_generation 

3 

4from .shared import _create_base_parser 

5 

6logger = bilby.core.utils.logger 

7 

8 

9def remove_argument_from_parser(parser, arg): 

10 for action in parser._actions: 

11 if action.dest == arg.replace("-", "_"): 

12 try: 

13 parser._handle_conflict_resolve(None, [("--" + arg, action)]) 

14 except ValueError as e: 

15 logger.warning(f"Error removing {arg}: {e}") 

16 logger.debug(f"Request to remove arg {arg} from bilby_pipe args, but arg not found") 

17 

18 

19def _add_slurm_settings_to_parser(parser): 

20 slurm_group = parser.add_argument_group(title="Slurm Settings") 

21 slurm_group.add_argument( 

22 "--nodes", type=int, default=1, help="Number of nodes to use (default 1)" 

23 ) 

24 slurm_group.add_argument( 

25 "--ntasks-per-node", 

26 type=int, 

27 default=2, 

28 help="Number of tasks per node (default 2)", 

29 ) 

30 slurm_group.add_argument( 

31 "--time", 

32 type=str, 

33 default="24:00:00", 

34 help="Maximum wall time (defaults to 24:00:00)", 

35 ) 

36 slurm_group.add_argument( 

37 "--mem-per-cpu", 

38 type=str, 

39 default="2G", 

40 help="Memory per CPU (defaults to 2GB)", 

41 ) 

42 slurm_group.add_argument( 

43 "--extra-lines", 

44 type=str, 

45 default=None, 

46 help="Additional lines, separated by ';', use for setting up conda env or module imports", 

47 ) 

48 slurm_group.add_argument( 

49 "--slurm-extra-lines", 

50 type=str, 

51 default=None, 

52 help="additional slurm args (args that need #SBATCH in front) of the form arg=val separated by sapce", 

53 ) 

54 return parser 

55 

56 

57def _create_reduced_bilby_pipe_parser(): 

58 bilby_pipe_parser = bilby_pipe.parser.create_parser() 

59 bilby_pipe_arguments_to_ignore = [ 

60 "version", 

61 "accounting", 

62 "local", 

63 "local-generation", 

64 "local-plot", 

65 "request-memory", 

66 "request-memory-generation", 

67 "request-cpus", 

68 "singularity-image", 

69 "scheduler", 

70 "scheduler-args", 

71 "scheduler-module", 

72 "scheduler-env", 

73 "transfer-files", 

74 "online-pe", 

75 "osg", 

76 "email", 

77 "postprocessing-executable", 

78 "postprocessing-arguments", 

79 "sampler", 

80 "sampling-seed", 

81 "sampler-kwargs", 

82 "sampler_kwargs", 

83 "plot-calibration", 

84 "plot-corner", 

85 "plot-format", 

86 "plot-marginal", 

87 "plot-skymap", 

88 "plot-waveform", 

89 ] 

90 for arg in bilby_pipe_arguments_to_ignore: 

91 remove_argument_from_parser(bilby_pipe_parser, arg) 

92 

93 bilby_pipe_parser.add_argument( 

94 "--sampler", 

95 choices=["dynesty"], 

96 default="dynesty", 

97 type=str, 

98 help="The parallelised sampler to use, defaults to dynesty", 

99 ) 

100 return bilby_pipe_parser 

101 

102 

103def add_extra_args_from_bilby_pipe_namespace(cli_args, parallel_bilby_args): 

104 """ 

105 :param args: args from parallel_bilby 

106 :return: Namespace argument object 

107 """ 

108 pipe_args, _ = bilby_pipe.data_generation.parse_args( 

109 cli_args, bilby_pipe.data_generation.create_generation_parser() 

110 ) 

111 for key, val in vars(pipe_args).items(): 

112 if key not in parallel_bilby_args: 

113 setattr(parallel_bilby_args, key, val) 

114 return parallel_bilby_args 

115 

116 

117def create_generation_parser(): 

118 """Parser for parallel_bilby_generation""" 

119 parser = _create_base_parser(sampler="all") 

120 bilby_pipe_parser = _create_reduced_bilby_pipe_parser() 

121 generation_parser = bilby_pipe.parser.BilbyArgParser( 

122 prog="parallel_bilby_generation", 

123 usage=__doc__, 

124 ignore_unknown_config_file_keys=False, 

125 allow_abbrev=False, 

126 parents=[parser, bilby_pipe_parser], 

127 add_help=False, 

128 ) 

129 generation_parser = _add_slurm_settings_to_parser(generation_parser) 

130 return generation_parser 

131 

132 

133def parse_generation_args(parser, cli_args=[""], as_namespace=False): 

134 """ 

135 Returns dictionary of arguments, as specified in the 

136 parser. 

137 

138 If no cli_args arguments are specified, returns the default arguments 

139 (by running the parser with no ini file and no CLI arguments) 

140 

141 Parameters 

142 ---------- 

143 parser: generation-parser 

144 cli_args: list of strings (default: [""]) 

145 List of arguments to be parsed. If empty, returns default arguments 

146 as_namespace: bool (default False) 

147 If True, returns the arguments as a Namespace object. If False, returns a dict 

148 

149 Returns 

150 ------- 

151 args: dict or Namespace 

152 

153 """ 

154 

155 args = parser.parse_args(args=cli_args) 

156 args = add_extra_args_from_bilby_pipe_namespace(cli_args, args) 

157 if as_namespace: 

158 return args 

159 return vars(args)