Coverage for python/lsst/ctrl/bps/pre_transform.py : 15%

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/>.
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
32import time
34from lsst.daf.butler import DimensionUniverse
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 (%s)", input_qgraph_filename)
77 stime = time.time()
78 qgraph_filename = os.path.join(out_prefix, os.path.basename(input_qgraph_filename))
79 shutil.copy2(input_qgraph_filename, qgraph_filename)
80 _LOG.info("Copying quantum graph took %.2f seconds", time.time() - stime)
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 (%s)", user_exec_butler_dir)
94 stime = time.time()
95 shutil.copytree(user_exec_butler_dir, execution_butler_dir)
96 _LOG.info("Copying execution butler took %.2f seconds", time.time() - stime)
97 else:
98 if when_create.upper() == "USER_PROVIDED":
99 raise KeyError("Missing qgraphFile to go with provided executionButlerDir")
101 # Run command to create the QuantumGraph.
102 _LOG.info("Creating quantum graph")
103 stime = time.time()
104 qgraph_filename = create_quantum_graph(config, out_prefix)
105 _LOG.info("Creating quantum graph took %.2f seconds", time.time() - stime)
107 _LOG.info("Reading quantum graph (%s)", qgraph_filename)
108 stime = time.time()
109 qgraph = read_quantum_graph(qgraph_filename)
110 _LOG.info("Reading quantum graph with %d nodes took %.2f seconds", len(qgraph),
111 time.time() - stime)
113 if when_create.upper() == "QGRAPH_CMDLINE":
114 if not os.path.exists(execution_butler_dir):
115 raise OSError(f"Missing execution butler dir ({execution_butler_dir}) after "
116 f"creating QuantumGraph (whenMakeExecutionButler == QGRAPH_CMDLINE")
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,
144 stderr=subprocess.STDOUT
145 )
146 buffer = os.read(process.stdout.fileno(), buffer_size).decode()
147 while process.poll is None or buffer:
148 print(buffer, end="", file=fh)
149 _LOG.info(buffer)
150 buffer = os.read(process.stdout.fileno(), buffer_size).decode()
151 process.stdout.close()
152 process.wait()
153 return process.returncode
156def create_quantum_graph(config, out_prefix=""):
157 """Create QuantumGraph from pipeline definition.
159 Parameters
160 ----------
161 config : `lsst.ctrl.bps.BpsConfig`
162 BPS configuration.
163 out_prefix : `str`, optional
164 Path in which to output QuantumGraph as well as the stdout/stderr
165 from generating the QuantumGraph. Defaults to empty string so
166 code will write the QuantumGraph and stdout/stderr to the current
167 directory.
169 Returns
170 -------
171 qgraph_filename : `str`
172 Name of file containing generated QuantumGraph.
173 """
174 # Create name of file to store QuantumGraph.
175 qgraph_filename = os.path.join(out_prefix, config['qgraphFileTemplate'])
177 # Get QuantumGraph generation command.
178 search_opt = {"curvals": {"qgraphFile": qgraph_filename}}
179 found, cmd = config.search("createQuantumGraph", opt=search_opt)
180 if not found:
181 _LOG.error("command for generating QuantumGraph not found")
182 _LOG.info(cmd)
184 # Run QuantumGraph generation.
185 out = os.path.join(out_prefix, "quantumGraphGeneration.out")
186 status = execute(cmd, out)
187 if status != 0:
188 raise RuntimeError(f"QuantumGraph generation exited with non-zero exit code ({status})\n"
189 f"Check {out} for more details.")
190 return qgraph_filename
193def read_quantum_graph(qgraph_filename):
194 """Read the QuantumGraph from disk.
196 Parameters
197 ----------
198 qgraph_filename : `str`
199 Name of file containing QuantumGraph to be used for workflow
200 generation.
202 Returns
203 -------
204 qgraph : `lsst.pipe.base.graph.QuantumGraph`
205 The QuantumGraph read from a file.
207 Raises
208 ------
209 RuntimeError
210 If the QuantumGraph contains 0 Quanta.
211 """
212 qgraph = QuantumGraph.loadUri(qgraph_filename, DimensionUniverse())
213 if len(qgraph) == 0:
214 raise RuntimeError("QuantumGraph is empty")
215 return qgraph
218def cluster_quanta(config, qgraph, name):
219 """Call specified function to group quanta into clusters to be run
220 together.
222 Parameters
223 ----------
224 config : `lsst.ctrl.bps.BpsConfig`
225 BPS configuration.
226 qgraph : `lsst.pipe.base.QuantumGraph`
227 Original full QuantumGraph for the run.
228 name : `str`
229 Name for the ClusteredQuantumGraph that will be generated.
231 Returns
232 -------
233 graph : `lsst.ctrl.bps.ClusteredQuantumGraph`
234 Generated ClusteredQuantumGraph.
235 """
236 cluster_func = doImport(config["clusterAlgorithm"])
237 return cluster_func(config, qgraph, name)