Coverage for python/lsst/ctrl/bps/bps_utils.py: 27%
83 statements
« prev ^ index » next coverage.py v7.3.1, created at 2023-10-02 08:09 +0000
« prev ^ index » next coverage.py v7.3.1, created at 2023-10-02 08:09 +0000
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 software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <https://www.gnu.org/licenses/>.
28"""Misc supporting classes and functions for BPS.
29"""
31__all__ = [
32 "chdir",
33 "create_job_quantum_graph_filename",
34 "save_qg_subgraph",
35 "_create_execution_butler",
36 "create_count_summary",
37 "parse_count_summary",
38 "_dump_pkg_info",
39 "_dump_env_info",
40]
42import contextlib
43import dataclasses
44import logging
45import os
46import shlex
47import subprocess
48from collections import Counter
49from enum import Enum
50from pathlib import Path
52import yaml
53from lsst.utils.packages import Packages
55_LOG = logging.getLogger(__name__)
58class WhenToSaveQuantumGraphs(Enum):
59 """Values for when to save the job quantum graphs."""
61 QGRAPH = 1 # Must be using single_quantum_clustering algorithm.
62 TRANSFORM = 2
63 PREPARE = 3
64 SUBMIT = 4
65 NEVER = 5 # Always use full QuantumGraph.
68@contextlib.contextmanager
69def chdir(path):
70 """Change working directory.
72 A chdir function that can be used inside a context.
74 Parameters
75 ----------
76 path : `str` or `pathlib.Path`
77 Path to be made current working directory.
78 """
79 cur_dir = os.getcwd()
80 os.chdir(path)
81 try:
82 yield
83 finally:
84 os.chdir(cur_dir)
87def create_job_quantum_graph_filename(config, job, out_prefix=None):
88 """Create a filename to be used when storing the QuantumGraph for a job.
90 Parameters
91 ----------
92 config : `lsst.ctrl.bps.BpsConfig`
93 BPS configuration.
94 job : `lsst.ctrl.bps.GenericWorkflowJob`
95 Job for which the QuantumGraph file is being saved.
96 out_prefix : `str`, optional
97 Path prefix for the QuantumGraph filename. If no out_prefix is given,
98 uses current working directory.
100 Returns
101 -------
102 full_filename : `str`
103 The filename for the job's QuantumGraph.
104 """
105 curvals = dataclasses.asdict(job)
106 if job.tags:
107 curvals.update(job.tags)
108 found, subdir = config.search("subDirTemplate", opt={"curvals": curvals})
109 if not found:
110 subdir = "{job.label}"
111 full_filename = Path("inputs") / subdir / f"quantum_{job.name}.qgraph"
113 if out_prefix is not None:
114 full_filename = Path(out_prefix) / full_filename
116 return str(full_filename)
119def save_qg_subgraph(qgraph, out_filename, node_ids=None):
120 """Save subgraph to file.
122 Parameters
123 ----------
124 qgraph : `lsst.pipe.base.QuantumGraph`
125 QuantumGraph to save.
126 out_filename : `str`
127 Name of the output file.
128 node_ids : `list` [`lsst.pipe.base.NodeId`]
129 NodeIds for the subgraph to save to file.
130 """
131 if not os.path.exists(out_filename):
132 _LOG.debug("Saving QuantumGraph with %d nodes to %s", len(qgraph), out_filename)
133 if node_ids is None:
134 qgraph.saveUri(out_filename)
135 else:
136 qgraph.subset(qgraph.getQuantumNodeByNodeId(nid) for nid in node_ids).saveUri(out_filename)
137 else:
138 _LOG.debug("Skipping saving QuantumGraph to %s because already exists.", out_filename)
141def _create_execution_butler(config, qgraph_filename, execution_butler_dir, out_prefix):
142 """Create the execution butler for use by the compute jobs.
144 Parameters
145 ----------
146 config : `lsst.ctrl.bps.BpsConfig`
147 BPS configuration.
148 qgraph_filename : `str`
149 Run QuantumGraph filename.
150 execution_butler_dir : `str`
151 Directory in which to create the execution butler.
152 out_prefix : `str` or None
153 Prefix for output filename to contain both stdout and stderr.
155 Raises
156 ------
157 CalledProcessError
158 Raised if command to create execution butler exits with non-zero
159 exit code.
160 """
161 _, command = config.search(
162 ".executionButler.createCommand",
163 opt={
164 "curvals": {"executionButlerDir": execution_butler_dir, "qgraphFile": qgraph_filename},
165 "replaceVars": True,
166 },
167 )
168 out_filename = "execution_butler_creation.out"
169 if out_prefix is not None:
170 out_filename = os.path.join(out_prefix, out_filename)
172 # When creating the execution Butler, handle separately errors related
173 # to creating the log file and errors directly related to creating
174 # the execution Butler itself.
175 opening = "cannot create the execution Butler"
176 try:
177 with open(out_filename, "w", encoding="utf-8") as fh:
178 print(command, file=fh)
179 print("\n", file=fh) # Note: want a blank line
180 subprocess.run(shlex.split(command), shell=False, check=True, stdout=fh, stderr=subprocess.STDOUT)
181 except OSError as exc:
182 raise type(exc)(f"{opening}: {exc.strerror}") from None
183 except subprocess.SubprocessError as exc:
184 raise RuntimeError(f"{opening}, see '{out_filename}' for details") from exc
187def create_count_summary(counts):
188 """Create summary from count mapping.
190 Parameters
191 ----------
192 counts : `collections.Counter` or `dict` [`str`, `int`]
193 Mapping of counts to keys.
195 Returns
196 -------
197 summary : `str`
198 Semi-colon delimited string of key:count pairs.
199 (e.g. "key1:cnt1;key2;cnt2") Parsable by
200 parse_count_summary().
201 """
202 summary = ""
203 if isinstance(counts, dict):
204 summary = ";".join([f"{key}:{counts[key]}" for key in counts])
205 return summary
208def parse_count_summary(summary):
209 """Parse summary into count mapping.
211 Parameters
212 ----------
213 summary : `str`
214 Semi-colon delimited string of key:count pairs.
216 Returns
217 -------
218 counts : `collections.Counter`
219 Mapping representation of given summary for easier
220 individual count lookup.
221 """
222 counts = Counter()
223 for part in summary.split(";"):
224 label, count = part.split(":")
225 counts[label] = count
226 return counts
229def _dump_pkg_info(filename):
230 """Save information about versions of packages in use for future reference.
232 Parameters
233 ----------
234 filename : `str`
235 The name of the file where to save the information about the versions
236 of the packages.
237 """
238 file = Path(filename)
239 if file.suffix.lower() not in {".yaml", ".yml"}:
240 file = file.with_suffix(f"{file.suffix}.yaml")
241 packages = Packages.fromSystem()
242 packages.write(str(file))
245def _dump_env_info(filename):
246 """Save information about runtime environment for future reference.
248 Parameters
249 ----------
250 filename : `str`
251 The name of the file where to save the information about the runtime
252 environment.
253 """
254 file = Path(filename)
255 if file.suffix.lower() not in {".yaml", ".yml"}:
256 file = file.with_suffix(f"{file.suffix}.yaml")
257 with open(file, "w", encoding="utf-8") as fh:
258 yaml.dump(dict(os.environ), fh)