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 

37import htcondor 

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 = {key.lower(): (type(val), val) for key, val in htcondor.param.items() if "DAGMAN_" in key} 

55 

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

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

58# 

59# Notes 

60# ----- 

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

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

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

64_fields.update( 

65 { 

66 "dagman_debug": (str, ""), 

67 "dagman_node_record_info": (str, ""), 

68 "dagman_record_machine_attrs": (str, ""), 

69 "dagman_manager_job_append_getenv": (str, ""), 

70 } 

71) 

72 

73# Dynamically create a Pydantic model encapsulating the DAGMan configuration 

74# options gathered above. 

75_DagmanOptions = create_model( 

76 "DagmanOptions", 

77 __config__=ConfigDict( 

78 alias_generator=AliasGenerator( 

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

80 ), 

81 extra="allow", 

82 serialize_by_alias=True, 

83 ), 

84 **_fields, 

85) 

86 

87 

88class DagmanConfigurator: 

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

90 

91 Parameters 

92 ---------- 

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

94 BPS configuration. 

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

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

97 

98 Raises 

99 ------ 

100 KeyError 

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

102 """ 

103 

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

105 if search_opts is None: 

106 search_opts = {} 

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

108 if site: 

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

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

111 if not wms_config: 

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

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

114 if self._options.model_extra: 

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

116 _LOG.warning( 

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

118 ", ".join(unknown_opts), 

119 ) 

120 self.config_path: Path | None = None 

121 self.prefix: Path | None = None 

122 

123 @property 

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

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

126 return { 

127 key: val 

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

129 if key not in self._options.model_extra 

130 } 

131 

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

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

134 

135 Parameters 

136 ---------- 

137 filename : `str`, optional 

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

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

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

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

142 

143 Raises 

144 ------ 

145 OSError 

146 Raised if the configuration file cannot be created. 

147 """ 

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

149 self.prefix = Path(prefix) 

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

151 try: 

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

153 except OSError as exc: 

154 _LOG.error( 

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

156 self.config_path, 

157 exc.strerror, 

158 ) 

159 raise 

160 

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

162 # explicitly set in the BPS configuration. 

163 # 

164 # Notes 

165 # ----- 

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

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

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

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

170 # manually. 

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

172 

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

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

175 

176 Parameters 

177 ---------- 

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

179 HTCondor DAG. 

180 

181 Raises 

182 ------ 

183 RuntimeError 

184 Raised if the prepare step was omitted. 

185 

186 Notes 

187 ----- 

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

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

190 """ 

191 if self.config_path is None: 

192 raise RuntimeError( 

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

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

195 ) 

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

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