Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1# This file is part of ctrl_bps. 

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 program is free software: you can redistribute it and/or modify 

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

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

12# (at your option) any later version. 

13# 

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

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

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

17# GNU General Public License for more details. 

18# 

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

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

21 

22"""Misc supporting classes and functions for BPS. 

23""" 

24 

25import dataclasses 

26import os 

27import shlex 

28import subprocess 

29import contextlib 

30import logging 

31from pathlib import Path 

32from enum import Enum 

33from collections import Counter 

34 

35 

36_LOG = logging.getLogger(__name__) 

37 

38 

39class WhenToSaveQuantumGraphs(Enum): 

40 """Values for when to save the job quantum graphs. 

41 """ 

42 QGRAPH = 1 # Must be using single_quantum_clustering algorithm. 

43 TRANSFORM = 2 

44 PREPARE = 3 

45 SUBMIT = 4 

46 NEVER = 5 # Always use full QuantumGraph. 

47 

48 

49@contextlib.contextmanager 

50def chdir(path): 

51 """A chdir function that can be used inside a context. 

52 

53 Parameters 

54 ---------- 

55 path : `str` 

56 Path to be made current working directory 

57 """ 

58 cur_dir = os.getcwd() 

59 os.chdir(path) 

60 try: 

61 yield 

62 finally: 

63 os.chdir(cur_dir) 

64 

65 

66def create_job_quantum_graph_filename(config, job, out_prefix=None): 

67 """Create a filename to be used when storing the QuantumGraph 

68 for a job. 

69 

70 Parameters 

71 ---------- 

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

73 BPS configuration. 

74 job : `lsst.ctrl.bps.GenericWorkflowJob` 

75 Job for which the QuantumGraph file is being saved. 

76 out_prefix : `str`, optional 

77 Path prefix for the QuantumGraph filename. If no out_prefix is given, 

78 uses current working directory. 

79 

80 Returns 

81 ------- 

82 full_filename : `str` 

83 The filename for the job's QuantumGraph. 

84 """ 

85 curvals = dataclasses.asdict(job) 

86 if job.tags: 

87 curvals.update(job.tags) 

88 found, subdir = config.search("subDirTemplate", opt={'curvals': curvals}) 

89 if not found: 

90 subdir = "{job.label}" 

91 full_filename = Path("inputs") / subdir / f"quantum_{job.name}.qgraph" 

92 

93 if out_prefix is not None: 

94 full_filename = Path(out_prefix) / full_filename 

95 

96 return str(full_filename) 

97 

98 

99def save_qg_subgraph(qgraph, out_filename, node_ids=None): 

100 """Save subgraph to file. 

101 

102 Parameters 

103 ---------- 

104 qgraph : `lsst.pipe.base.QuantumGraph` 

105 QuantumGraph to save. 

106 out_filename : `str` 

107 Name of the output file. 

108 node_ids : `list` [`lsst.pipe.base.NodeId`] 

109 NodeIds for the subgraph to save to file. 

110 """ 

111 if not os.path.exists(out_filename): 

112 _LOG.debug("Saving QuantumGraph with %d nodes to %s", len(qgraph), out_filename) 

113 if node_ids is None: 

114 qgraph.saveUri(out_filename) 

115 else: 

116 qgraph.subset(qgraph.getQuantumNodeByNodeId(nid) for nid in node_ids).saveUri(out_filename) 

117 else: 

118 _LOG.debug("Skipping saving QuantumGraph to %s because already exists.", out_filename) 

119 

120 

121def _create_execution_butler(config, qgraph_filename, execution_butler_dir, out_prefix): 

122 """Create the execution butler for use by the compute jobs. 

123 

124 Parameters 

125 ---------- 

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

127 BPS configuration. 

128 qgraph_filename : `str` 

129 Run QuantumGraph filename. 

130 execution_butler_dir : `str` 

131 Directory in which to create the execution butler. 

132 out_prefix : `str` or None 

133 Prefix for output filename to contain both stdout and stderr. 

134 

135 Raises 

136 ------ 

137 CalledProcessError 

138 Raised if command to create execution butler exits with non-zero 

139 exit code. 

140 """ 

141 _, command = config.search(".executionButler.createCommand", 

142 opt={"curvals": {"executionButlerDir": execution_butler_dir, 

143 "qgraphFile": qgraph_filename}, 

144 "replaceVars": True}) 

145 out_filename = "execution_butler_creation.out" 

146 if out_prefix is not None: 

147 out_filename = os.path.join(out_prefix, out_filename) 

148 with open(out_filename, "w") as fh: 

149 print(command, file=fh) 

150 print("\n", file=fh) # Note: want a blank line 

151 subprocess.run(shlex.split(command), shell=False, check=True, stdout=fh, stderr=subprocess.STDOUT) 

152 

153 

154def create_count_summary(counts): 

155 """Create summary from count mapping. 

156 

157 Parameters 

158 ---------- 

159 count : `collections.Counter` or `dict` [`str`, `int`] 

160 Mapping of counts to keys. 

161 

162 Returns 

163 ------- 

164 summary : `str` 

165 Semi-colon delimited string of key:count pairs. 

166 (e.g. "key1:cnt1;key2;cnt2") Parsable by 

167 parse_count_summary(). 

168 """ 

169 summary = "" 

170 if isinstance(counts, dict): 

171 summary = ";".join([f"{key}:{counts[key]}" for key in counts]) 

172 return summary 

173 

174 

175def parse_count_summary(summary): 

176 """Parse summary into count mapping. 

177 

178 Parameters 

179 ---------- 

180 summary : `str` 

181 Semi-colon delimited string of key:count pairs. 

182 

183 Returns 

184 ------- 

185 counts : `collections.Counter` 

186 Mapping representation of given summary for easier 

187 individual count lookup. 

188 """ 

189 counts = Counter() 

190 for part in summary.split(";"): 

191 label, count = part.split(":") 

192 counts[label] = count 

193 return counts