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

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
33from collections import Counter
36_LOG = logging.getLogger(__name__)
39class WhenToSaveQuantumGraphs(Enum):
40 """Values for when to save the job quantum graphs.
41 """
42 QGRAPH = 1 # Must be using single_quantum_clustering algorithm.
43 TRANSFORM = 2
44 PREPARE = 3
45 SUBMIT = 4
46 NEVER = 5 # Always use full QuantumGraph.
49@contextlib.contextmanager
50def chdir(path):
51 """A chdir function that can be used inside a context.
53 Parameters
54 ----------
55 path : `str`
56 Path to be made current working directory
57 """
58 cur_dir = os.getcwd()
59 os.chdir(path)
60 try:
61 yield
62 finally:
63 os.chdir(cur_dir)
66def create_job_quantum_graph_filename(config, job, out_prefix=None):
67 """Create a filename to be used when storing the QuantumGraph
68 for a job.
70 Parameters
71 ----------
72 config : `lsst.ctrl.bps.BpsConfig`
73 BPS configuration (at minimum must contain qgraphFile and
74 outCollection).
75 job : `lsst.ctrl.bps.GenericWorkflowJob`
76 Job for which the QuantumGraph file is being saved.
77 out_prefix : `str`, optional
78 Path prefix for the QuantumGraph filename. If no out_prefix is given,
79 uses current working directory.
81 Returns
82 -------
83 full_filename : `str`
84 The filename for the job's QuantumGraph.
85 """
86 curvals = dataclasses.asdict(job)
87 if job.tags:
88 curvals.update(job.tags)
89 found, subdir = config.search("subDirTemplate", opt={'curvals': curvals})
90 if not found:
91 subdir = "{job.label}"
92 full_filename = Path("inputs") / subdir / f"quantum_{job.name}.qgraph"
94 if out_prefix is not None:
95 full_filename = Path(out_prefix) / full_filename
97 return str(full_filename)
100def save_qg_subgraph(qgraph, out_filename, node_ids=None):
101 """Save subgraph to file.
103 Parameters
104 ----------
105 qgraph : `lsst.pipe.base.QuantumGraph`
106 QuantumGraph to save.
107 out_filename : `str`
108 Name of the output file.
109 node_ids : `list` [`lsst.pipe.base.NodeId`]
110 NodeIds for the subgraph to save to file.
111 """
112 if not os.path.exists(out_filename):
113 _LOG.debug("Saving QuantumGraph with %d nodes to %s", len(qgraph), out_filename)
114 if node_ids is None:
115 qgraph.saveUri(out_filename)
116 else:
117 qgraph.subset(qgraph.getQuantumNodeByNodeId(nid) for nid in node_ids).saveUri(out_filename)
118 else:
119 _LOG.debug("Skipping saving QuantumGraph to %s because already exists.", out_filename)
122def _create_execution_butler(config, qgraph_filename, execution_butler_dir, out_prefix):
123 """Create the execution butler for use by the compute jobs.
125 Parameters
126 ----------
127 config : `lsst.ctrl.bps.BpsConfig`
128 BPS configuration (at minimum must contain qgraphFile and
129 outCollection).
130 qgraph_filename : `str`
131 Run QuantumGraph filename.
132 execution_butler_dir : `str`
133 Directory in which to create the execution butler.
134 out_prefix : `str` or None
135 Prefix for output filename to contain both stdout and stderr.
137 Raises
138 ------
139 CalledProcessError
140 Raised if command to create execution butler exits with non-zero
141 exit code.
142 """
143 _, command = config.search(".executionButler.createCommand",
144 opt={"curvals": {"executionButlerDir": execution_butler_dir,
145 "qgraphFile": qgraph_filename},
146 "replaceVars": True})
147 out_filename = "execution_butler_creation.out"
148 if out_prefix is not None:
149 out_filename = os.path.join(out_prefix, out_filename)
150 with open(out_filename, "w") as fh:
151 print(command, file=fh)
152 print("\n", file=fh) # Note: want a blank line
153 subprocess.run(shlex.split(command), shell=False, check=True, stdout=fh, stderr=subprocess.STDOUT)
156def create_count_summary(counts):
157 """Create summary from count mapping.
159 Parameters
160 ----------
161 count : `collections.Counter` or `dict` [`str`, `int`]
162 Mapping of counts to keys.
164 Returns
165 -------
166 summary : `str`
167 Semi-colon delimited string of key:count pairs.
168 (e.g. "key1:cnt1;key2;cnt2") Parsable by
169 parse_count_summary().
170 """
171 summary = ""
172 if isinstance(counts, dict):
173 summary = ";".join([f"{key}:{counts[key]}" for key in counts])
174 return summary
177def parse_count_summary(summary):
178 """Parse summary into count mapping.
180 Parameters
181 ----------
182 summary : `str`
183 Semi-colon delimited string of key:count pairs.
185 Returns
186 -------
187 counts : `collections.Counter`
188 Mapping representation of given summary for easier
189 individual count lookup.
190 """
191 counts = Counter()
192 for part in summary.split(";"):
193 label, count = part.split(":")
194 counts[label] = count
195 return counts