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

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 os
26import contextlib
27import logging
28from enum import Enum
30_LOG = logging.getLogger()
33class WhenToSaveQuantumGraphs(Enum):
34 """Values for when to save the job quantum graphs."""
35 QGRAPH = 1 # Must be using single_quantum_clustering algorithm.
36 TRANSFORM = 2
37 PREPARE = 3
38 SUBMIT = 4
39 NEVER = 5 # Always use full QuantumGraph.
42@contextlib.contextmanager
43def chdir(path):
44 """A chdir function that can be used inside a context.
46 Parameters
47 ----------
48 path : `str`
49 Path to be made current working directory
50 """
51 cur_dir = os.getcwd()
52 os.chdir(path)
53 try:
54 yield
55 finally:
56 os.chdir(cur_dir)
59def create_job_quantum_graph_filename(job, out_prefix=None):
60 """Create a filename to be used when storing the QuantumGraph
61 for a job.
63 Parameters
64 ----------
65 job : `~lsst.ctrl.bps.generic_workflow.GenericWorkflowJob`
66 Job for which the QuantumGraph file is being saved.
67 out_prefix : `str`, optional
68 Path prefix for the QuantumGraph filename. If no
69 out_prefix is given, uses current working directory.
71 Returns
72 -------
73 full_filename : `str`
74 The filename for the job's QuantumGraph.
75 """
76 name_parts = []
77 if out_prefix is not None:
78 name_parts.append(out_prefix)
79 name_parts.append("inputs")
80 if job.label is not None:
81 name_parts.append(job.label)
82 name_parts.append(f"quantum_{job.name}.qgraph")
83 full_filename = os.path.join("", *name_parts)
84 return full_filename
87def save_qg_subgraph(qgraph, out_filename, node_ids=None):
88 """Save subgraph to file.
90 Parameters
91 ----------
92 qgraph : `~lsst.pipe.base.QuantumGraph`
93 QuantumGraph to save.
94 out_filename : `str`
95 Name of the output file.
96 node_ids : `list` of `~lsst.pipe.base.NodeId`
97 NodeIds for the subgraph to save to file.
98 """
99 if not os.path.exists(out_filename):
100 _LOG.debug("Saving QuantumGraph with %d nodes to %s", len(qgraph), out_filename)
101 if node_ids is None:
102 qgraph.saveUri(out_filename)
103 else:
104 qgraph.subset(qgraph.getQuantumNodeByNodeId(nid) for nid in node_ids).saveUri(out_filename)
105 else:
106 _LOG.debug("Skipping saving QuantumGraph to %s because already exists.", out_filename)