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 

33 

34 

35_LOG = logging.getLogger(__name__) 

36 

37 

38class WhenToSaveQuantumGraphs(Enum): 

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

40 """ 

41 QGRAPH = 1 # Must be using single_quantum_clustering algorithm. 

42 TRANSFORM = 2 

43 PREPARE = 3 

44 SUBMIT = 4 

45 NEVER = 5 # Always use full QuantumGraph. 

46 

47 

48@contextlib.contextmanager 

49def chdir(path): 

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

51 

52 Parameters 

53 ---------- 

54 path : `str` 

55 Path to be made current working directory 

56 """ 

57 cur_dir = os.getcwd() 

58 os.chdir(path) 

59 try: 

60 yield 

61 finally: 

62 os.chdir(cur_dir) 

63 

64 

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

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

67 for a job. 

68 

69 Parameters 

70 ---------- 

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

72 BPS configuration (at minimum must contain qgraphFile and 

73 outCollection). 

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 (at minimum must contain qgraphFile and 

128 outCollection). 

129 qgraph_filename : `str` 

130 Run QuantumGraph filename. 

131 execution_butler_dir : `str` 

132 Directory in which to create the execution butler. 

133 out_prefix : `str` or None 

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

135 

136 Raises 

137 ------ 

138 CalledProcessError 

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

140 exit code. 

141 """ 

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

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

144 "qgraphFile": qgraph_filename}, 

145 "replaceVars": True}) 

146 out_filename = "execution_butler_creation.out" 

147 if out_prefix is not None: 

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

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

150 print(command, file=fh) 

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

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