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

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 shutil
47from lsst.obs.base import Instrument
48from lsst.utils.timer import time_this
50from . import BPS_SEARCH_ORDER, BpsConfig
51from .pre_transform import acquire_quantum_graph, cluster_quanta
52from .transform import transform
53from .prepare import prepare
54from .submit import submit
55from .cancel import cancel
56from .report import report
59_LOG = logging.getLogger(__name__)
62def _init_submission_driver(config_file, **kwargs):
63 """Initialize runtime environment.
65 Parameters
66 ----------
67 config_file : `str`
68 Name of the configuration file.
70 Returns
71 -------
72 config : `lsst.ctrl.bps.BpsConfig`
73 Batch Processing Service configuration.
74 """
75 config = BpsConfig(config_file, BPS_SEARCH_ORDER)
77 # Override config with command-line values
78 # Handle diffs between pipetask argument names vs bps yaml
79 translation = {"input": "inCollection",
80 "output_run": "outputRun",
81 "qgraph": "qgraphFile",
82 "pipeline": "pipelineYaml"}
83 for key, value in kwargs.items():
84 # Don't want to override config with None or empty string values.
85 if value:
86 # pipetask argument parser converts some values to list,
87 # but bps will want string.
88 if not isinstance(value, str):
89 value = ",".join(value)
90 new_key = translation.get(key, re.sub(r"_(\S)", lambda match: match.group(1).upper(), key))
91 config[f".bps_cmdline.{new_key}"] = value
93 # Set some initial values
94 config[".bps_defined.timestamp"] = Instrument.makeCollectionTimestamp()
95 if "operator" not in config:
96 config[".bps_defined.operator"] = getpass.getuser()
98 if "outCollection" in config:
99 raise KeyError("outCollection is deprecated. Replace all outCollection references with outputRun.")
101 if "outputRun" not in config:
102 raise KeyError("Must specify the output run collection using outputRun")
104 if "uniqProcName" not in config:
105 config[".bps_defined.uniqProcName"] = config["outputRun"].replace("/", "_")
107 if "submitPath" not in config:
108 raise KeyError("Must specify the submit-side run directory using submitPath")
110 # make submit directory to contain all outputs
111 submit_path = config["submitPath"]
112 os.makedirs(submit_path, exist_ok=True)
113 config[".bps_defined.submitPath"] = submit_path
115 # save copy of configs (orig and expanded config)
116 shutil.copy2(config_file, submit_path)
117 with open(f"{submit_path}/{config['uniqProcName']}_config.yaml", "w") as fh:
118 config.dump(fh)
120 return config
123def acquire_qgraph_driver(config_file, **kwargs):
124 """Read a quantum graph from a file or create one from pipeline definition.
126 Parameters
127 ----------
128 config_file : `str`
129 Name of the configuration file.
131 Returns
132 -------
133 config : `lsst.ctrl.bps.BpsConfig`
134 Updated configuration.
135 qgraph : `lsst.pipe.base.graph.QuantumGraph`
136 A graph representing quanta.
137 """
138 config = _init_submission_driver(config_file, **kwargs)
139 submit_path = config[".bps_defined.submitPath"]
141 _LOG.info("Starting acquire stage (generating and/or reading quantum graph)")
142 with time_this(log=_LOG, level=logging.INFO, prefix=None, msg="Acquire stage completed"):
143 qgraph_file, qgraph, execution_butler_dir = acquire_quantum_graph(config, out_prefix=submit_path)
145 config[".bps_defined.executionButlerDir"] = execution_butler_dir
146 config[".bps_defined.runQgraphFile"] = qgraph_file
147 return config, qgraph
150def cluster_qgraph_driver(config_file, **kwargs):
151 """Group quanta into clusters.
153 Parameters
154 ----------
155 config_file : `str`
156 Name of the configuration file.
158 Returns
159 -------
160 config : `lsst.ctrl.bps.BpsConfig`
161 Updated configuration.
162 clustered_qgraph : `lsst.ctrl.bps.ClusteredQuantumGraph`
163 A graph representing clustered quanta.
164 """
165 config, qgraph = acquire_qgraph_driver(config_file, **kwargs)
167 _LOG.info("Starting cluster stage (grouping quanta into jobs)")
168 with time_this(log=_LOG, level=logging.INFO, prefix=None, msg="Cluster stage completed"):
169 clustered_qgraph = cluster_quanta(config, qgraph, config["uniqProcName"])
171 submit_path = config[".bps_defined.submitPath"]
172 _, save_clustered_qgraph = config.search("saveClusteredQgraph", opt={"default": False})
173 if save_clustered_qgraph:
174 clustered_qgraph.save(os.path.join(submit_path, "bps_clustered_qgraph.pickle"))
175 _, save_dot = config.search("saveDot", opt={"default": False})
176 if save_dot:
177 clustered_qgraph.draw(os.path.join(submit_path, "bps_clustered_qgraph.dot"))
178 return config, clustered_qgraph
181def transform_driver(config_file, **kwargs):
182 """Create a workflow for a specific workflow management system.
184 Parameters
185 ----------
186 config_file : `str`
187 Name of the configuration file.
189 Returns
190 -------
191 generic_workflow_config : `lsst.ctrl.bps.BpsConfig`
192 Configuration to use when creating the workflow.
193 generic_workflow : `lsst.ctrl.bps.BaseWmsWorkflow`
194 Representation of the abstract/scientific workflow specific to a given
195 workflow management system.
196 """
197 config, clustered_qgraph = cluster_qgraph_driver(config_file, **kwargs)
198 submit_path = config[".bps_defined.submitPath"]
200 _LOG.info("Starting transform stage (creating generic workflow)")
201 with time_this(log=_LOG, level=logging.INFO, prefix=None, msg="Transform stage completed"):
202 generic_workflow, generic_workflow_config = transform(config, clustered_qgraph, submit_path)
203 _LOG.info("Generic workflow name '%s'", generic_workflow.name)
205 _, save_workflow = config.search("saveGenericWorkflow", opt={"default": False})
206 if save_workflow:
207 with open(os.path.join(submit_path, "bps_generic_workflow.pickle"), "wb") as outfh:
208 generic_workflow.save(outfh, "pickle")
209 _, save_dot = config.search("saveDot", opt={"default": False})
210 if save_dot:
211 with open(os.path.join(submit_path, "bps_generic_workflow.dot"), "w") as outfh:
212 generic_workflow.draw(outfh, "dot")
213 return generic_workflow_config, generic_workflow
216def prepare_driver(config_file, **kwargs):
217 """Create a representation of the generic workflow.
219 Parameters
220 ----------
221 config_file : `str`
222 Name of the configuration file.
224 Returns
225 -------
226 wms_config : `lsst.ctrl.bps.BpsConfig`
227 Configuration to use when creating the workflow.
228 workflow : `lsst.ctrl.bps.BaseWmsWorkflow`
229 Representation of the abstract/scientific workflow specific to a given
230 workflow management system.
231 """
232 generic_workflow_config, generic_workflow = transform_driver(config_file, **kwargs)
233 submit_path = generic_workflow_config[".bps_defined.submitPath"]
235 _LOG.info("Starting prepare stage (creating specific implementation of workflow)")
236 with time_this(log=_LOG, level=logging.INFO, prefix=None, msg="Prepare stage completed"):
237 wms_workflow = prepare(generic_workflow_config, generic_workflow, submit_path)
239 wms_workflow_config = generic_workflow_config
240 print(f"Submit dir: {wms_workflow.submit_path}")
241 return wms_workflow_config, wms_workflow
244def submit_driver(config_file, **kwargs):
245 """Submit workflow for execution.
247 Parameters
248 ----------
249 config_file : `str`
250 Name of the configuration file.
251 """
252 _LOG.info("Starting submission process")
253 with time_this(log=_LOG, level=logging.INFO, prefix=None, msg="Completed entire submission process"):
254 wms_workflow_config, wms_workflow = prepare_driver(config_file, **kwargs)
256 _LOG.info("Starting submit stage")
257 with time_this(log=_LOG, level=logging.INFO, prefix=None, msg="Completed submit stage"):
258 submit(wms_workflow_config, wms_workflow)
259 _LOG.info("Run '%s' submitted for execution with id '%s'", wms_workflow.name, wms_workflow.run_id)
261 print(f"Run Id: {wms_workflow.run_id}")
264def report_driver(wms_service, run_id, user, hist_days, pass_thru):
265 """Print out summary of jobs submitted for execution.
267 Parameters
268 ----------
269 wms_service : `str`
270 Name of the class.
271 run_id : `str`
272 A run id the report will be restricted to.
273 user : `str`
274 A user name the report will be restricted to.
275 hist_days : int
276 Number of days
277 pass_thru : `str`
278 A string to pass directly to the WMS service class.
279 """
280 report(wms_service, run_id, user, hist_days, pass_thru)
283def cancel_driver(wms_service, run_id, user, require_bps, pass_thru):
284 """Cancel submitted workflows.
286 Parameters
287 ----------
288 wms_service : `str`
289 Name of the Workload Management System service class.
290 run_id : `str`
291 ID or path of job that should be canceled.
292 user : `str`
293 User whose submitted jobs should be canceled.
294 require_bps : `bool`
295 Whether to require given run_id/user to be a bps submitted job.
296 pass_thru : `str`
297 Information to pass through to WMS.
298 """
299 cancel(wms_service, run_id, user, require_bps, pass_thru)