Coverage for python/lsst/ctrl/bps/bps_utils.py : 33%

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"""Misc supporting classes and functions for BPS.
23"""
25import dataclasses
26import os
27import shlex
28import subprocess
29import contextlib
30import logging
31from pathlib import Path
32from enum import Enum
35_LOG = logging.getLogger(__name__)
38class WhenToSaveQuantumGraphs(Enum):
39 """Values for when to save the job quantum graphs.
40 """
41 QGRAPH = 1 # Must be using single_quantum_clustering algorithm.
42 TRANSFORM = 2
43 PREPARE = 3
44 SUBMIT = 4
45 NEVER = 5 # Always use full QuantumGraph.
48@contextlib.contextmanager
49def chdir(path):
50 """A chdir function that can be used inside a context.
52 Parameters
53 ----------
54 path : `str`
55 Path to be made current working directory
56 """
57 cur_dir = os.getcwd()
58 os.chdir(path)
59 try:
60 yield
61 finally:
62 os.chdir(cur_dir)
65def create_job_quantum_graph_filename(config, job, out_prefix=None):
66 """Create a filename to be used when storing the QuantumGraph
67 for a job.
69 Parameters
70 ----------
71 config : `lsst.ctrl.bps.BpsConfig`
72 BPS configuration (at minimum must contain qgraphFile and
73 outCollection).
74 job : `lsst.ctrl.bps.GenericWorkflowJob`
75 Job for which the QuantumGraph file is being saved.
76 out_prefix : `str`, optional
77 Path prefix for the QuantumGraph filename. If no out_prefix is given,
78 uses current working directory.
80 Returns
81 -------
82 full_filename : `str`
83 The filename for the job's QuantumGraph.
84 """
85 curvals = dataclasses.asdict(job)
86 if job.tags:
87 curvals.update(job.tags)
88 found, subdir = config.search("subDirTemplate", opt={'curvals': curvals})
89 if not found:
90 subdir = "{job.label}"
91 full_filename = Path("inputs") / subdir / f"quantum_{job.name}.qgraph"
93 if out_prefix is not None:
94 full_filename = Path(out_prefix) / full_filename
96 return str(full_filename)
99def save_qg_subgraph(qgraph, out_filename, node_ids=None):
100 """Save subgraph to file.
102 Parameters
103 ----------
104 qgraph : `lsst.pipe.base.QuantumGraph`
105 QuantumGraph to save.
106 out_filename : `str`
107 Name of the output file.
108 node_ids : `list` [`lsst.pipe.base.NodeId`]
109 NodeIds for the subgraph to save to file.
110 """
111 if not os.path.exists(out_filename):
112 _LOG.debug("Saving QuantumGraph with %d nodes to %s", len(qgraph), out_filename)
113 if node_ids is None:
114 qgraph.saveUri(out_filename)
115 else:
116 qgraph.subset(qgraph.getQuantumNodeByNodeId(nid) for nid in node_ids).saveUri(out_filename)
117 else:
118 _LOG.debug("Skipping saving QuantumGraph to %s because already exists.", out_filename)
121def _create_execution_butler(config, qgraph_filename, execution_butler_dir, out_prefix):
122 """Create the execution butler for use by the compute jobs.
124 Parameters
125 ----------
126 config : `lsst.ctrl.bps.BpsConfig`
127 BPS configuration (at minimum must contain qgraphFile and
128 outCollection).
129 qgraph_filename : `str`
130 Run QuantumGraph filename.
131 execution_butler_dir : `str`
132 Directory in which to create the execution butler.
133 out_prefix : `str` or None
134 Prefix for output filename to contain both stdout and stderr.
136 Raises
137 ------
138 CalledProcessError
139 Raised if command to create execution butler exits with non-zero
140 exit code.
141 """
142 _, command = config.search(".executionButler.createCommand",
143 opt={"curvals": {"executionButlerDir": execution_butler_dir,
144 "qgraphFile": qgraph_filename},
145 "replaceVars": True})
146 out_filename = "execution_butler_creation.out"
147 if out_prefix is not None:
148 out_filename = os.path.join(out_prefix, out_filename)
149 with open(out_filename, "w") as fh:
150 print(command, file=fh)
151 print("\n", file=fh) # Note: want a blank line
152 subprocess.run(shlex.split(command), shell=False, check=True, stdout=fh, stderr=subprocess.STDOUT)