Coverage for python/lsst/ctrl/bps/pre_transform.py: 15%
76 statements
« prev ^ index » next coverage.py v7.2.4, created at 2023-04-29 03:27 -0700
« prev ^ index » next coverage.py v7.2.4, created at 2023-04-29 03:27 -0700
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.pipe.base.graph import QuantumGraph
35from lsst.utils import doImport
36from lsst.utils.logging import VERBOSE
37from lsst.utils.timer import time_this, timeMethod
39_LOG = logging.getLogger(__name__)
42@timeMethod(logger=_LOG, logLevel=VERBOSE)
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(
94 log=_LOG, level=logging.INFO, prefix=None, msg="Completed copying execution butler"
95 ):
96 shutil.copytree(user_exec_butler_dir, execution_butler_dir)
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 with time_this(log=_LOG, level=logging.INFO, prefix=None, msg="Completed creating quantum graph"):
104 qgraph_filename = create_quantum_graph(config, out_prefix)
106 _LOG.info("Reading quantum graph from '%s'", qgraph_filename)
107 with time_this(log=_LOG, level=logging.INFO, prefix=None, msg="Completed reading quantum graph"):
108 qgraph = QuantumGraph.loadUri(qgraph_filename)
110 if when_create.upper() == "QGRAPH_CMDLINE":
111 if not os.path.exists(execution_butler_dir):
112 raise OSError(
113 f"Missing execution butler dir ({execution_butler_dir}) after "
114 "creating QuantumGraph (whenMakeExecutionButler == QGRAPH_CMDLINE"
115 )
116 elif when_create.upper() == "ACQUIRE":
117 _create_execution_butler(config, qgraph_filename, execution_butler_dir, config["submitPath"])
119 return qgraph_filename, qgraph, execution_butler_dir
122def execute(command, filename):
123 """Execute a command.
125 Parameters
126 ----------
127 command : `str`
128 String representing the command to execute.
129 filename : `str`
130 A file to which both stderr and stdout will be written to.
132 Returns
133 -------
134 exit_code : `int`
135 The exit code the command being executed finished with.
136 """
137 buffer_size = 5000
138 with open(filename, "w") as fh:
139 print(command, file=fh)
140 print("\n", file=fh) # Note: want a blank line
141 process = subprocess.Popen(
142 shlex.split(command), shell=False, stdout=subprocess.PIPE, stderr=subprocess.STDOUT
143 )
144 buffer = os.read(process.stdout.fileno(), buffer_size).decode()
145 while process.poll is None or buffer:
146 print(buffer, end="", file=fh)
147 _LOG.info(buffer)
148 buffer = os.read(process.stdout.fileno(), buffer_size).decode()
149 process.stdout.close()
150 process.wait()
151 return process.returncode
154def create_quantum_graph(config, out_prefix=""):
155 """Create QuantumGraph from pipeline definition.
157 Parameters
158 ----------
159 config : `lsst.ctrl.bps.BpsConfig`
160 BPS configuration.
161 out_prefix : `str`, optional
162 Path in which to output QuantumGraph as well as the stdout/stderr
163 from generating the QuantumGraph. Defaults to empty string so
164 code will write the QuantumGraph and stdout/stderr to the current
165 directory.
167 Returns
168 -------
169 qgraph_filename : `str`
170 Name of file containing generated QuantumGraph.
171 """
172 # Create name of file to store QuantumGraph.
173 qgraph_filename = os.path.join(out_prefix, config["qgraphFileTemplate"])
175 # Get QuantumGraph generation command.
176 search_opt = {"curvals": {"qgraphFile": qgraph_filename}}
177 found, cmd = config.search("createQuantumGraph", opt=search_opt)
178 if not found:
179 _LOG.error("command for generating QuantumGraph not found")
180 _LOG.info(cmd)
182 # Run QuantumGraph generation.
183 out = os.path.join(out_prefix, "quantumGraphGeneration.out")
184 status = execute(cmd, out)
185 if status != 0:
186 raise RuntimeError(
187 f"QuantumGraph generation exited with non-zero exit code ({status})\n"
188 f"Check {out} for more details."
189 )
190 return qgraph_filename
193def cluster_quanta(config, qgraph, name):
194 """Call specified function to group quanta into clusters to be run
195 together.
197 Parameters
198 ----------
199 config : `lsst.ctrl.bps.BpsConfig`
200 BPS configuration.
201 qgraph : `lsst.pipe.base.QuantumGraph`
202 Original full QuantumGraph for the run.
203 name : `str`
204 Name for the ClusteredQuantumGraph that will be generated.
206 Returns
207 -------
208 graph : `lsst.ctrl.bps.ClusteredQuantumGraph`
209 Generated ClusteredQuantumGraph.
210 """
211 cluster_func = doImport(config["clusterAlgorithm"])
212 return cluster_func(config, qgraph, name)