Coverage for python/lsst/ctrl/bps/drivers.py : 22%

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 <http://www.gnu.org/licenses/>.
22"""Driver functions for each subcommand.
24Driver functions ensure that ensure all setup work is done before running
25the subcommand method.
26"""
29__all__ = [
30 "acquire_qgraph_driver",
31 "cluster_qgraph_driver",
32 "transform_driver",
33 "prepare_driver",
34 "submit_driver",
35 "report_driver",
36 "cancel_driver",
37]
40import getpass
41import logging
42import os
43import re
44import pickle
45import time
46import shutil
49from lsst.obs.base import Instrument
51from . import BPS_SEARCH_ORDER, BpsConfig
52from .bps_draw import draw_networkx_dot
53from .pre_transform import acquire_quantum_graph, cluster_quanta
54from .transform import transform
55from .prepare import prepare
56from .submit import submit
57from .cancel import cancel
58from .report import report
61_LOG = logging.getLogger(__name__)
64def _init_submission_driver(config_file, **kwargs):
65 """Initialize runtime environment.
67 Parameters
68 ----------
69 config_file : `str`
70 Name of the configuration file.
72 Returns
73 -------
74 config : `lsst.ctrl.bps.BpsConfig`
75 Batch Processing Service configuration.
76 """
77 config = BpsConfig(config_file, BPS_SEARCH_ORDER)
79 # Override config with command-line values
80 # Handle diffs between pipetask argument names vs bps yaml
81 translation = {"input": "inCollection",
82 "output_run": "outCollection",
83 "qgraph": "qgraphFile",
84 "pipeline": "pipelineYaml"}
85 for key, value in kwargs.items():
86 # Don't want to override config with None or empty string values.
87 if value:
88 # pipetask argument parser converts some values to list,
89 # but bps will want string.
90 if not isinstance(value, str):
91 value = ",".join(value)
92 new_key = translation.get(key, re.sub(r"_(\S)", lambda match: match.group(1).upper(), key))
93 config[f".bps_cmdline.{new_key}"] = value
95 # Set some initial values
96 config[".bps_defined.timestamp"] = Instrument.makeCollectionTimestamp()
97 if "operator" not in config:
98 config[".bps_defined.operator"] = getpass.getuser()
100 if "uniqProcName" not in config:
101 config[".bps_defined.uniqProcName"] = config["outCollection"].replace("/", "_")
103 # make submit directory to contain all outputs
104 submit_path = config["submitPath"]
105 os.makedirs(submit_path, exist_ok=True)
106 config[".bps_defined.submitPath"] = submit_path
108 # save copy of configs (orig and expanded config)
109 shutil.copy2(config_file, submit_path)
110 with open(f"{submit_path}/{config['uniqProcName']}_config.yaml", "w") as fh:
111 config.dump(fh)
113 return config
116def acquire_qgraph_driver(config_file, **kwargs):
117 """Read a quantum graph from a file or create one from pipeline definition.
119 Parameters
120 ----------
121 config_file : `str`
122 Name of the configuration file.
124 Returns
125 -------
126 config : `lsst.ctrl.bps.BpsConfig`
127 Updated configuration.
128 qgraph : `lsst.pipe.base.graph.QuantumGraph`
129 A graph representing quanta.
130 """
131 stime = time.time()
132 config = _init_submission_driver(config_file, **kwargs)
133 submit_path = config[".bps_defined.submitPath"]
134 _LOG.info("Acquiring QuantumGraph (it will be created from pipeline definition if needed)")
135 qgraph_file, qgraph, execution_butler_dir = acquire_quantum_graph(config, out_prefix=submit_path)
136 config[".bps_defined.executionButlerDir"] = execution_butler_dir
137 _LOG.info("Run QuantumGraph file %s", qgraph_file)
138 config[".bps_defined.runQgraphFile"] = qgraph_file
139 _LOG.info("Acquiring QuantumGraph took %.2f seconds", time.time() - stime)
141 return config, qgraph
144def cluster_qgraph_driver(config_file, **kwargs):
145 """Group quanta into clusters.
147 Parameters
148 ----------
149 config_file : `str`
150 Name of the configuration file.
152 Returns
153 -------
154 config : `lsst.ctrl.bps.BpsConfig`
155 Updated configuration.
156 clustered_qgraph : `lsst.ctrl.bps.ClusteredQuantumGraph`
157 A graph representing clustered quanta.
158 """
159 stime = time.time()
160 config, qgraph = acquire_qgraph_driver(config_file, **kwargs)
161 _LOG.info("Clustering quanta")
162 clustered_qgraph = cluster_quanta(config, qgraph, config["uniqProcName"])
163 _LOG.info("Clustering quanta took %.2f seconds", time.time() - stime)
165 submit_path = config[".bps_defined.submitPath"]
166 _, save_clustered_qgraph = config.search("saveClusteredQgraph", opt={"default": False})
167 if save_clustered_qgraph:
168 with open(os.path.join(submit_path, "bps_clustered_qgraph.pickle"), "wb") as outfh:
169 pickle.dump(clustered_qgraph, outfh)
170 _, save_dot = config.search("saveDot", opt={"default": False})
171 if save_dot:
172 draw_networkx_dot(clustered_qgraph, os.path.join(submit_path, "bps_clustered_qgraph.dot"))
173 return config, clustered_qgraph
176def transform_driver(config_file, **kwargs):
177 """Create a workflow for a specific workflow management system.
179 Parameters
180 ----------
181 config_file : `str`
182 Name of the configuration file.
184 Returns
185 -------
186 generic_workflow_config : `lsst.ctrl.bps.BpsConfig`
187 Configuration to use when creating the workflow.
188 generic_workflow : `lsst.ctrl.bps.BaseWmsWorkflow`
189 Representation of the abstract/scientific workflow specific to a given
190 workflow management system.
191 """
192 stime = time.time()
193 config, clustered_qgraph = cluster_qgraph_driver(config_file, **kwargs)
194 submit_path = config[".bps_defined.submitPath"]
195 _LOG.info("Creating Generic Workflow")
196 generic_workflow, generic_workflow_config = transform(config, clustered_qgraph, submit_path)
197 _LOG.info("Creating Generic Workflow took %.2f seconds", time.time() - stime)
198 _LOG.info("Generic Workflow name %s", generic_workflow.name)
200 _, save_workflow = config.search("saveGenericWorkflow", opt={"default": False})
201 if save_workflow:
202 with open(os.path.join(submit_path, "bps_generic_workflow.pickle"), "wb") as outfh:
203 generic_workflow.save(outfh, "pickle")
204 _, save_dot = config.search("saveDot", opt={"default": False})
205 if save_dot:
206 with open(os.path.join(submit_path, "bps_generic_workflow.dot"), "w") as outfh:
207 generic_workflow.draw(outfh, "dot")
208 return generic_workflow_config, generic_workflow
211def prepare_driver(config_file, **kwargs):
212 """Create a representation of the generic workflow.
214 Parameters
215 ----------
216 config_file : `str`
217 Name of the configuration file.
219 Returns
220 -------
221 wms_config : `lsst.ctrl.bps.BpsConfig`
222 Configuration to use when creating the workflow.
223 workflow : `lsst.ctrl.bps.BaseWmsWorkflow`
224 Representation of the abstract/scientific workflow specific to a given
225 workflow management system.
226 """
227 stime = time.time()
228 generic_workflow_config, generic_workflow = transform_driver(config_file, **kwargs)
229 submit_path = generic_workflow_config[".bps_defined.submitPath"]
230 _LOG.info("Creating specific implementation of workflow")
231 wms_workflow = prepare(generic_workflow_config, generic_workflow, submit_path)
232 wms_workflow_config = generic_workflow_config
233 _LOG.info("Creating specific implementation of workflow took %.2f seconds", time.time() - stime)
234 print(f"Submit dir: {wms_workflow.submit_path}")
235 return wms_workflow_config, wms_workflow
238def submit_driver(config_file, **kwargs):
239 """Submit workflow for execution.
241 Parameters
242 ----------
243 config_file : `str`
244 Name of the configuration file.
245 """
246 wms_workflow_config, wms_workflow = prepare_driver(config_file, **kwargs)
247 submit(wms_workflow_config, wms_workflow)
248 print(f"Run Id: {wms_workflow.run_id}")
251def report_driver(wms_service, run_id, user, hist_days, pass_thru):
252 """Print out summary of jobs submitted for execution.
254 Parameters
255 ----------
256 wms_service : `str`
257 Name of the class.
258 run_id : `str`
259 A run id the report will be restricted to.
260 user : `str`
261 A user name the report will be restricted to.
262 hist_days : int
263 Number of days
264 pass_thru : `str`
265 A string to pass directly to the WMS service class.
266 """
267 report(wms_service, run_id, user, hist_days, pass_thru)
270def cancel_driver(wms_service, run_id, user, require_bps, pass_thru):
271 """Cancel submitted workflows.
273 Parameters
274 ----------
275 wms_service : `str`
276 Name of the Workload Management System service class.
277 run_id : `str`
278 ID or path of job that should be canceled.
279 user : `str`
280 User whose submitted jobs should be canceled.
281 require_bps : `bool`
282 Whether to require given run_id/user to be a bps submitted job.
283 pass_thru : `str`
284 Information to pass through to WMS.
285 """
286 cancel(wms_service, run_id, user, require_bps, pass_thru)