Coverage for python/lsst/ctrl/bps/htcondor/dagman_configurator.py: 98%

1# This file is part of ctrl_bps_htcondor. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (https://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# This software is dual licensed under the GNU General Public License and also 

10# under a 3-clause BSD license. Recipients may choose which of these licenses 

11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt, 

12# respectively. If you choose the GPL option then the following text applies 

13# (but note that there is still no warranty even if you opt for BSD instead): 

14# 

15# This program is free software: you can redistribute it and/or modify 

16# it under the terms of the GNU General Public License as published by 

17# the Free Software Foundation, either version 3 of the License, or 

18# (at your option) any later version. 

19# 

20# This program is distributed in the hope that it will be useful, 

21# but WITHOUT ANY WARRANTY; without even the implied warranty of 

22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

23# GNU General Public License for more details. 

24# 

25# You should have received a copy of the GNU General Public License 

26# along with this program. If not, see <https://www.gnu.org/licenses/>. 

27 

28"""Module enabling configuring DAGMan via submit YAML.""" 

29 

30__all__ = ["DagmanConfigurator"] 

31 

32import logging 

33import os 

34from pathlib import Path 

35from typing import Any 

36 

37from htcondor2 import param as htc_param 

38from pydantic import AliasGenerator, ConfigDict, create_model 

39 

40from lsst.ctrl.bps import BpsConfig 

41 

42from .lssthtc import HTCDag 

43 

44_LOG = logging.getLogger(__name__) 

45 

46# Extract DAGMan configuration options with their types and default values from 

47# the local HTCondor configuration. 

48# 

49# Notes 

50# ----- 

51# There are some DAGMan configuration options that names do not start with 

52# ``DAGMAN_`` (e.g., ``MAX_DAGMAN_LOG``). Hence, do not use 

53# ``key.startswith("DAGMAN_")``. 

54_fields: dict[str, Any] = { 

55 key.lower(): (type(val), val) for key, val in htc_param.items() if "DAGMAN_" in key 

56} 

57 

58# Add some valid configuration options are not set by default by HTCondor and 

59# are missing from ``htcondor.param``. 

60# 

61# Notes 

62# ----- 

63# A complete list of configuration options HTCondor supports can be found in 

64# ``src/condor_utils/param_info.in`` in 

65# `HTCondor GitHub repository <https://github.com/htcondor/htcondor>`_. 

66_fields.update( 

67 { 

68 "dagman_config_file": (str, ""), 

69 "dagman_debug": (str, ""), 

70 "dagman_insert_sub_file": (str, ""), 

71 "dagman_node_record_info": (str, ""), 

72 "dagman_record_machine_attrs": (str, ""), 

73 } 

74) 

75 

76# Dynamically create a Pydantic model encapsulating the DAGMan configuration 

77# options gathered above. 

78_DagmanOptions = create_model( 

79 "DagmanOptions", 

80 __config__=ConfigDict( 

81 alias_generator=AliasGenerator( 

82 serialization_alias=lambda name: name.upper(), 

83 ), 

84 extra="allow", 

85 serialize_by_alias=True, 

86 ), 

87 **_fields, 

88) 

89 

90 

91class DagmanConfigurator: 

92 """Class responsible for setting WMS-specific configuration options. 

93 

94 Parameters 

95 ---------- 

96 config : `lsst.ctrl.bps.BpsConfig` 

97 BPS configuration. 

98 search_opts : `dict` [`str`, `~typing.Any`], optional 

99 Options to use while searching the BPS configuration for values. 

100 

101 Raises 

102 ------ 

103 KeyError 

104 Raised if DAGMan configuration is missing from the BPS configuration. 

105 """ 

106 

107 def __init__(self, config: BpsConfig, search_opts: dict[str, Any] | None = None) -> None: 

108 if search_opts is None: 

109 search_opts = {} 

110 _, site = config.search("computeSite", search_opts) 

111 if site: 

112 search_opts["curvals"] = {"curr_site": site} 

113 _, wms_config = config.search("wmsConfig", search_opts) 

114 if not wms_config: 

115 raise KeyError("WMS-specific configuration not found") 

116 self._options = _DagmanOptions.model_validate({key.lower(): val for key, val in wms_config.items()}) 

117 if self._options.model_extra: 

118 unknown_opts = [key.upper() for key in self._options.model_extra] 

119 _LOG.warning( 

120 "The following WMS-specific config options were not recognized and will be ignored: %s.", 

121 ", ".join(unknown_opts), 

122 ) 

123 self.config_path: Path | None = None 

124 self.prefix: Path | None = None 

125 

126 @property 

127 def options(self) -> dict[str, Any]: 

128 """DAGMan configuration options set via BPS (`dict` [`str`, `Any`]).""" 

129 return { 

130 key: val 

131 for key, val in self._options.model_dump(exclude_unset=True).items() 

132 if key not in self._options.model_extra 

133 } 

134 

135 def prepare(self, filename: os.PathLike | str, prefix: os.PathLike | str | None) -> None: 

136 """Write WMS-specific configuration to a file. 

137 

138 Parameters 

139 ---------- 

140 filename : `str`, optional 

141 Name of the file to use when creating the DAG configuration. 

142 prefix : `pathlib.Path` | `str`, optional 

143 Directory in which to output the DAG configuration file. If not 

144 provided, the script will be written to the current directory. 

145 

146 Raises 

147 ------ 

148 OSError 

149 Raised if the configuration file cannot be created. 

150 """ 

151 if prefix: 151 ↛ 153line 151 didn't jump to line 153 because the condition on line 151 was always true

152 self.prefix = Path(prefix) 

153 self.config_path = self.prefix / filename if self.prefix else Path(filename) 

154 try: 

155 self.config_path.parent.mkdir(parents=True, exist_ok=True) 

156 except OSError as exc: 

157 _LOG.error( 

158 "Could not write WMS-specific configuration file '%s': %s", 

159 self.config_path, 

160 exc.strerror, 

161 ) 

162 raise 

163 

164 # Populate the DAG configuration file only with options that were 

165 # explicitly set in the BPS configuration. 

166 # 

167 # Notes 

168 # ----- 

169 # The Pydantic model we are using to represent the DAGMan configuration 

170 # options allows for extra fields. However, it seems that 

171 # BaseModel.model_dump() does not support excluding these fields during 

172 # serialization at the moment (Pydantic ver. 2.12), so we have to do it 

173 # manually. 

174 self.config_path.write_text("\n".join(f"{key} = {val}" for key, val in self.options.items())) 

175 

176 def configure(self, dag: HTCDag) -> None: 

177 """Add DAG configuration file to the workflow. 

178 

179 Parameters 

180 ---------- 

181 dag : `lsst.ctrl.bps.htcondor.HTCDag` 

182 HTCondor DAG. 

183 

184 Raises 

185 ------ 

186 RuntimeError 

187 Raised if the prepare step was omitted. 

188 

189 Notes 

190 ----- 

191 The path to the DAG configuration is added as a DAG attribute named 

192 ``bps_wms_config_path``. The stored path is relative to the prefix. 

193 """ 

194 if self.config_path is None: 

195 raise RuntimeError( 

196 f"cannot add WMS-specific configuration to the workflow: file does not exist. " 

197 f"(hint: run {type(self).__qualname__}.prepare() to create it)" 

198 ) 

199 config_path = self.config_path.relative_to(self.prefix) if self.prefix else self.config_path 

200 dag.add_attribs({"bps_wms_config_path": str(config_path)})