Coverage for python/lsst/ctrl/bps/pre_transform.py: 16%
Shortcuts 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
Shortcuts 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/>.
22"""Driver to execute steps outside of BPS that need to be done first
23including running QuantumGraph generation and reading the QuantumGraph
24into memory.
25"""
27import logging
28import os
29import shlex
30import shutil
31import subprocess
33from lsst.daf.butler import DimensionUniverse
34from lsst.daf.butler.core.utils import time_this
35from lsst.pipe.base.graph import QuantumGraph
36from lsst.utils import doImport
37from lsst.ctrl.bps.bps_utils import _create_execution_butler
40_LOG = logging.getLogger(__name__)
43def acquire_quantum_graph(config, out_prefix=""):
44 """Read a quantum graph from a file or create one from scratch.
46 Parameters
47 ----------
48 config : `lsst.ctrl.bps.BpsConfig`
49 Configuration values for BPS. In particular, looking for qgraphFile.
50 out_prefix : `str`, optional
51 Output path for the QuantumGraph and stdout/stderr from generating
52 the QuantumGraph. Default value is empty string.
54 Returns
55 -------
56 qgraph_filename : `str`
57 Name of file containing QuantumGraph that was read into qgraph.
58 qgraph : `lsst.pipe.base.graph.QuantumGraph`
59 A QuantumGraph read in from pre-generated file or one that is the
60 result of running code that generates it.
61 execution_butler_dir : `str` or None
62 The directory containing the execution butler if user-provided or
63 created during this submission step.
64 """
65 # consistently name execution butler directory
66 _, execution_butler_dir = config.search("executionButlerTemplate")
67 if not execution_butler_dir.startswith('/'):
68 execution_butler_dir = os.path.join(config["submitPath"], execution_butler_dir)
69 _, when_create = config.search(".executionButler.whenCreate")
71 # Check to see if user provided pre-generated QuantumGraph.
72 found, input_qgraph_filename = config.search("qgraphFile")
73 if found and input_qgraph_filename:
74 if out_prefix is not None:
75 # Save a copy of the QuantumGraph file in out_prefix.
76 _LOG.info("Copying quantum graph from '%s'", input_qgraph_filename)
77 with time_this(log=_LOG, level=logging.INFO, prefix=None, msg="Completed copying quantum graph"):
78 qgraph_filename = os.path.join(out_prefix, os.path.basename(input_qgraph_filename))
79 shutil.copy2(input_qgraph_filename, qgraph_filename)
80 else:
81 # Use QuantumGraph file in original given location.
82 qgraph_filename = input_qgraph_filename
84 # Copy Execution Butler if user provided (shouldn't provide execution
85 # butler if not providing QuantumGraph)
86 if when_create.upper() == "USER_PROVIDED":
87 found, user_exec_butler_dir = config.search(".executionButler.executionButlerDir")
88 if not found:
89 raise KeyError("Missing .executionButler.executionButlerDir for when_create == USER_PROVIDED")
91 # Save a copy of the execution butler file in out_prefix.
92 _LOG.info("Copying execution butler to '%s'", user_exec_butler_dir)
93 with time_this(log=_LOG, level=logging.INFO, prefix=None,
94 msg="Completed copying execution butler"):
95 shutil.copytree(user_exec_butler_dir, execution_butler_dir)
96 else:
97 if when_create.upper() == "USER_PROVIDED":
98 raise KeyError("Missing qgraphFile to go with provided executionButlerDir")
100 # Run command to create the QuantumGraph.
101 _LOG.info("Creating quantum graph")
102 with time_this(log=_LOG, level=logging.INFO, prefix=None, msg="Completed creating quantum graph"):
103 qgraph_filename = create_quantum_graph(config, out_prefix)
105 _LOG.info("Reading quantum graph from '%s'", qgraph_filename)
106 with time_this(log=_LOG, level=logging.INFO, prefix=None, msg="Completed reading quantum graph"):
107 qgraph = read_quantum_graph(qgraph_filename)
109 if when_create.upper() == "QGRAPH_CMDLINE":
110 if not os.path.exists(execution_butler_dir):
111 raise OSError(f"Missing execution butler dir ({execution_butler_dir}) after "
112 f"creating QuantumGraph (whenMakeExecutionButler == QGRAPH_CMDLINE")
113 elif when_create.upper() == "ACQUIRE":
114 _create_execution_butler(config, qgraph_filename, execution_butler_dir, config["submitPath"])
116 return qgraph_filename, qgraph, execution_butler_dir
119def execute(command, filename):
120 """Execute a command.
122 Parameters
123 ----------
124 command : `str`
125 String representing the command to execute.
126 filename : `str`
127 A file to which both stderr and stdout will be written to.
129 Returns
130 -------
131 exit_code : `int`
132 The exit code the command being executed finished with.
133 """
134 buffer_size = 5000
135 with open(filename, "w") as fh:
136 print(command, file=fh)
137 print("\n", file=fh) # Note: want a blank line
138 process = subprocess.Popen(
139 shlex.split(command), shell=False, stdout=subprocess.PIPE,
140 stderr=subprocess.STDOUT
141 )
142 buffer = os.read(process.stdout.fileno(), buffer_size).decode()
143 while process.poll is None or buffer:
144 print(buffer, end="", file=fh)
145 _LOG.info(buffer)
146 buffer = os.read(process.stdout.fileno(), buffer_size).decode()
147 process.stdout.close()
148 process.wait()
149 return process.returncode
152def create_quantum_graph(config, out_prefix=""):
153 """Create QuantumGraph from pipeline definition.
155 Parameters
156 ----------
157 config : `lsst.ctrl.bps.BpsConfig`
158 BPS configuration.
159 out_prefix : `str`, optional
160 Path in which to output QuantumGraph as well as the stdout/stderr
161 from generating the QuantumGraph. Defaults to empty string so
162 code will write the QuantumGraph and stdout/stderr to the current
163 directory.
165 Returns
166 -------
167 qgraph_filename : `str`
168 Name of file containing generated QuantumGraph.
169 """
170 # Create name of file to store QuantumGraph.
171 qgraph_filename = os.path.join(out_prefix, config['qgraphFileTemplate'])
173 # Get QuantumGraph generation command.
174 search_opt = {"curvals": {"qgraphFile": qgraph_filename}}
175 found, cmd = config.search("createQuantumGraph", opt=search_opt)
176 if not found:
177 _LOG.error("command for generating QuantumGraph not found")
178 _LOG.info(cmd)
180 # Run QuantumGraph generation.
181 out = os.path.join(out_prefix, "quantumGraphGeneration.out")
182 status = execute(cmd, out)
183 if status != 0:
184 raise RuntimeError(f"QuantumGraph generation exited with non-zero exit code ({status})\n"
185 f"Check {out} for more details.")
186 return qgraph_filename
189def read_quantum_graph(qgraph_filename):
190 """Read the QuantumGraph from disk.
192 Parameters
193 ----------
194 qgraph_filename : `str`
195 Name of file containing QuantumGraph to be used for workflow
196 generation.
198 Returns
199 -------
200 qgraph : `lsst.pipe.base.graph.QuantumGraph`
201 The QuantumGraph read from a file.
203 Raises
204 ------
205 RuntimeError
206 If the QuantumGraph contains 0 Quanta.
207 """
208 qgraph = QuantumGraph.loadUri(qgraph_filename, DimensionUniverse())
209 if len(qgraph) == 0:
210 raise RuntimeError("QuantumGraph is empty")
211 return qgraph
214def cluster_quanta(config, qgraph, name):
215 """Call specified function to group quanta into clusters to be run
216 together.
218 Parameters
219 ----------
220 config : `lsst.ctrl.bps.BpsConfig`
221 BPS configuration.
222 qgraph : `lsst.pipe.base.QuantumGraph`
223 Original full QuantumGraph for the run.
224 name : `str`
225 Name for the ClusteredQuantumGraph that will be generated.
227 Returns
228 -------
229 graph : `lsst.ctrl.bps.ClusteredQuantumGraph`
230 Generated ClusteredQuantumGraph.
231 """
232 cluster_func = doImport(config["clusterAlgorithm"])
233 return cluster_func(config, qgraph, name)