Coverage for python/lsst/ctrl/bps/pre_transform.py: 16%
83 statements
« prev ^ index » next coverage.py v6.5.0, created at 2023-01-26 02:12 -0800
« prev ^ index » next coverage.py v6.5.0, created at 2023-01-26 02:12 -0800
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.ctrl.bps.bps_utils import _create_execution_butler
34from lsst.daf.butler import Butler
35from lsst.pipe.base.graph import QuantumGraph
36from lsst.utils import doImport
37from lsst.utils.logging import VERBOSE
38from lsst.utils.timer import time_this, timeMethod
40_LOG = logging.getLogger(__name__)
43@timeMethod(logger=_LOG, logLevel=VERBOSE)
44def acquire_quantum_graph(config, out_prefix=""):
45 """Read a quantum graph from a file or create one from scratch.
47 Parameters
48 ----------
49 config : `lsst.ctrl.bps.BpsConfig`
50 Configuration values for BPS. In particular, looking for qgraphFile.
51 out_prefix : `str`, optional
52 Output path for the QuantumGraph and stdout/stderr from generating
53 the QuantumGraph. Default value is empty string.
55 Returns
56 -------
57 qgraph_filename : `str`
58 Name of file containing QuantumGraph that was read into qgraph.
59 qgraph : `lsst.pipe.base.graph.QuantumGraph`
60 A QuantumGraph read in from pre-generated file or one that is the
61 result of running code that generates it.
62 execution_butler_dir : `str` or None
63 The directory containing the execution butler if user-provided or
64 created during this submission step.
65 """
66 # consistently name execution butler directory
67 _, execution_butler_dir = config.search("executionButlerTemplate")
68 if not execution_butler_dir.startswith("/"):
69 execution_butler_dir = os.path.join(config["submitPath"], execution_butler_dir)
70 _, when_create = config.search(".executionButler.whenCreate")
72 # Check to see if user provided pre-generated QuantumGraph.
73 found, input_qgraph_filename = config.search("qgraphFile")
74 if found and input_qgraph_filename:
75 if out_prefix is not None:
76 # Save a copy of the QuantumGraph file in out_prefix.
77 _LOG.info("Copying quantum graph from '%s'", input_qgraph_filename)
78 with time_this(log=_LOG, level=logging.INFO, prefix=None, msg="Completed copying quantum graph"):
79 qgraph_filename = os.path.join(out_prefix, os.path.basename(input_qgraph_filename))
80 shutil.copy2(input_qgraph_filename, qgraph_filename)
81 else:
82 # Use QuantumGraph file in original given location.
83 qgraph_filename = input_qgraph_filename
85 # Copy Execution Butler if user provided (shouldn't provide execution
86 # butler if not providing QuantumGraph)
87 if when_create.upper() == "USER_PROVIDED":
88 found, user_exec_butler_dir = config.search(".executionButler.executionButlerDir")
89 if not found:
90 raise KeyError("Missing .executionButler.executionButlerDir for when_create == USER_PROVIDED")
92 # Save a copy of the execution butler file in out_prefix.
93 _LOG.info("Copying execution butler to '%s'", user_exec_butler_dir)
94 with time_this(
95 log=_LOG, level=logging.INFO, prefix=None, msg="Completed copying execution butler"
96 ):
97 shutil.copytree(user_exec_butler_dir, execution_butler_dir)
98 else:
99 if when_create.upper() == "USER_PROVIDED":
100 raise KeyError("Missing qgraphFile to go with provided executionButlerDir")
102 # Run command to create the QuantumGraph.
103 _LOG.info("Creating quantum graph")
104 with time_this(log=_LOG, level=logging.INFO, prefix=None, msg="Completed creating quantum graph"):
105 qgraph_filename = create_quantum_graph(config, out_prefix)
107 _LOG.info("Reading quantum graph from '%s'", qgraph_filename)
108 with time_this(log=_LOG, level=logging.INFO, prefix=None, msg="Completed reading quantum graph"):
109 qgraph = read_quantum_graph(qgraph_filename, config["butlerConfig"])
111 if when_create.upper() == "QGRAPH_CMDLINE":
112 if not os.path.exists(execution_butler_dir):
113 raise OSError(
114 f"Missing execution butler dir ({execution_butler_dir}) after "
115 f"creating QuantumGraph (whenMakeExecutionButler == QGRAPH_CMDLINE"
116 )
117 elif when_create.upper() == "ACQUIRE":
118 _create_execution_butler(config, qgraph_filename, execution_butler_dir, config["submitPath"])
120 return qgraph_filename, qgraph, execution_butler_dir
123def execute(command, filename):
124 """Execute a command.
126 Parameters
127 ----------
128 command : `str`
129 String representing the command to execute.
130 filename : `str`
131 A file to which both stderr and stdout will be written to.
133 Returns
134 -------
135 exit_code : `int`
136 The exit code the command being executed finished with.
137 """
138 buffer_size = 5000
139 with open(filename, "w") as fh:
140 print(command, file=fh)
141 print("\n", file=fh) # Note: want a blank line
142 process = subprocess.Popen(
143 shlex.split(command), shell=False, stdout=subprocess.PIPE, stderr=subprocess.STDOUT
144 )
145 buffer = os.read(process.stdout.fileno(), buffer_size).decode()
146 while process.poll is None or buffer:
147 print(buffer, end="", file=fh)
148 _LOG.info(buffer)
149 buffer = os.read(process.stdout.fileno(), buffer_size).decode()
150 process.stdout.close()
151 process.wait()
152 return process.returncode
155def create_quantum_graph(config, out_prefix=""):
156 """Create QuantumGraph from pipeline definition.
158 Parameters
159 ----------
160 config : `lsst.ctrl.bps.BpsConfig`
161 BPS configuration.
162 out_prefix : `str`, optional
163 Path in which to output QuantumGraph as well as the stdout/stderr
164 from generating the QuantumGraph. Defaults to empty string so
165 code will write the QuantumGraph and stdout/stderr to the current
166 directory.
168 Returns
169 -------
170 qgraph_filename : `str`
171 Name of file containing generated QuantumGraph.
172 """
173 # Create name of file to store QuantumGraph.
174 qgraph_filename = os.path.join(out_prefix, config["qgraphFileTemplate"])
176 # Get QuantumGraph generation command.
177 search_opt = {"curvals": {"qgraphFile": qgraph_filename}}
178 found, cmd = config.search("createQuantumGraph", opt=search_opt)
179 if not found:
180 _LOG.error("command for generating QuantumGraph not found")
181 _LOG.info(cmd)
183 # Run QuantumGraph generation.
184 out = os.path.join(out_prefix, "quantumGraphGeneration.out")
185 status = execute(cmd, out)
186 if status != 0:
187 raise RuntimeError(
188 f"QuantumGraph generation exited with non-zero exit code ({status})\n"
189 f"Check {out} for more details."
190 )
191 return qgraph_filename
194def read_quantum_graph(qgraph_filename, butler_uri):
195 """Read the QuantumGraph from disk.
197 Parameters
198 ----------
199 qgraph_filename : `str`
200 Name of file containing QuantumGraph to be used for workflow
201 generation.
202 butler_uri : `str`
203 Location of butler repository that can be used to create a
204 butler object.
206 Returns
207 -------
208 qgraph : `lsst.pipe.base.graph.QuantumGraph`
209 The QuantumGraph read from a file.
211 Raises
212 ------
213 RuntimeError
214 If the QuantumGraph contains 0 Quanta.
215 """
216 # Get the DimensionUniverse from the butler repository
217 butler = Butler(butler_uri, writeable=False)
218 qgraph = QuantumGraph.loadUri(qgraph_filename, butler.registry.dimensions)
219 if len(qgraph) == 0:
220 raise RuntimeError("QuantumGraph is empty")
221 return qgraph
224def cluster_quanta(config, qgraph, name):
225 """Call specified function to group quanta into clusters to be run
226 together.
228 Parameters
229 ----------
230 config : `lsst.ctrl.bps.BpsConfig`
231 BPS configuration.
232 qgraph : `lsst.pipe.base.QuantumGraph`
233 Original full QuantumGraph for the run.
234 name : `str`
235 Name for the ClusteredQuantumGraph that will be generated.
237 Returns
238 -------
239 graph : `lsst.ctrl.bps.ClusteredQuantumGraph`
240 Generated ClusteredQuantumGraph.
241 """
242 cluster_func = doImport(config["clusterAlgorithm"])
243 return cluster_func(config, qgraph, name)