Coverage for tests/test_graph_walker.py: 17%
47 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-10-09 03:38 -0700
1# This file is part of pipe_base.
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/>.
28from __future__ import annotations
30import unittest
31import uuid
33import networkx
35from lsst.daf.butler import DataCoordinate
36from lsst.pipe.base.graph_walker import GraphWalker
37from lsst.pipe.base.tests.mocks import DynamicConnectionConfig, InMemoryRepo
40class GraphWalkerTestCase(unittest.TestCase):
41 """Tests for the GraphWalker utility class."""
43 def test_iteration(self) -> None:
44 helper = InMemoryRepo("base.yaml", "spatial.yaml")
45 helper.add_task("a", dimensions=["visit", "detector"])
46 helper.add_task("b", dimensions=["visit", "detector"])
47 helper.add_task(
48 "c", # Gathers outputs from 'b' by visit.
49 dimensions=["visit"],
50 inputs={
51 "input_connection": DynamicConnectionConfig(
52 dataset_type_name=f"dataset_auto{helper.last_auto_dataset_type_index}",
53 dimensions=["visit", "detector"],
54 multiple=True,
55 )
56 },
57 )
58 helper.add_task(
59 "d", # Scatters visit-level outputs from 'c' per-detector.
60 dimensions=["visit", "detector"],
61 inputs={
62 "input_connection": DynamicConnectionConfig(
63 dataset_type_name=f"dataset_auto{helper.last_auto_dataset_type_index}",
64 dimensions=["visit"],
65 )
66 },
67 )
68 qg = helper.make_quantum_graph_builder().finish(attach_datastore_records=False).assemble()
69 walker = GraphWalker[uuid.UUID](qg.quantum_only_xgraph.copy())
70 # First set of unblocked nodes consists of all 'a' nodes.
71 a_nodes: dict[DataCoordinate, uuid.UUID] = {
72 qg.quantum_only_xgraph.nodes[n]["data_id"]: n for n in next(walker)
73 }
74 self.assertEqual([qg.quantum_only_xgraph.nodes[n]["task_label"] for n in a_nodes.values()], ["a"] * 8)
75 # Mark the detector=1 'a' nodes as finished, and iterate the walker;
76 # this should unblock all of the detector=1 'b' nodes.
77 walker.finish(
78 a_nodes.pop(
79 DataCoordinate.standardize(
80 instrument="Cam1", visit=1, detector=1, universe=helper.butler.dimensions
81 )
82 )
83 )
84 walker.finish(
85 a_nodes.pop(
86 DataCoordinate.standardize(
87 instrument="Cam1", visit=2, detector=1, universe=helper.butler.dimensions
88 )
89 )
90 )
91 b_nodes: dict[DataCoordinate, uuid.UUID] = {
92 qg.quantum_only_xgraph.nodes[n]["data_id"]: n for n in next(walker)
93 }
94 self.assertEqual([qg.quantum_only_xgraph.nodes[n]["task_label"] for n in b_nodes.values()], ["b"] * 2)
95 self.assertEqual([data_id["detector"] for data_id in b_nodes.keys()], [1] * 2)
96 # Mark one quantum of 'a' as a failure, and check that this returns all
97 # downstream nodes.
98 bad_node = a_nodes.pop(
99 DataCoordinate.standardize(
100 instrument="Cam1", visit=1, detector=2, universe=helper.butler.dimensions
101 )
102 )
103 downstream_of_bad = set(walker.fail(bad_node))
104 self.assertEqual(downstream_of_bad, set(networkx.dag.descendants(qg.quantum_only_xgraph, bad_node)))
105 # Mark the remaining 'a' nodes as finished. This should unblock the
106 # remaining 5 'b' nodes, for a total of 7 unblocked.
107 for n in a_nodes.values():
108 walker.finish(n)
109 a_nodes.clear()
110 b_nodes.update({qg.quantum_only_xgraph.nodes[n]["data_id"]: n for n in next(walker)})
111 self.assertEqual([qg.quantum_only_xgraph.nodes[n]["task_label"] for n in b_nodes.values()], ["b"] * 7)
112 # Check that iterating the walker again yields nothing new. Note that
113 # this is not the same as the iterator being exhausted, which only
114 # happens after we've traversed the whole graph.
115 self.assertFalse(next(walker))
116 # Mark those 'b' nodes as finished, which should unblock one 'c' node
117 # (the other was downstream of the failure).
118 for n in b_nodes.values():
119 walker.finish(n)
120 c_nodes: dict[DataCoordinate, uuid.UUID] = {
121 qg.quantum_only_xgraph.nodes[n]["data_id"]: n for n in next(walker)
122 }
123 self.assertEqual(len(c_nodes), 1)
124 c_data_id, c_node = c_nodes.popitem()
125 self.assertEqual(qg.quantum_only_xgraph.nodes[c_node]["task_label"], "c")
126 self.assertEqual(
127 c_data_id,
128 DataCoordinate.standardize(instrument="Cam1", visit=2, universe=helper.butler.dimensions),
129 )
130 # Mark the 'c' node as finished, which should unblock the 4 'd' nodes
131 # for visit=2.
132 walker.finish(c_node)
133 d_nodes: dict[DataCoordinate, uuid.UUID] = {
134 qg.quantum_only_xgraph.nodes[n]["data_id"]: n for n in next(walker)
135 }
136 self.assertEqual([qg.quantum_only_xgraph.nodes[n]["task_label"] for n in d_nodes.values()], ["d"] * 4)
137 self.assertEqual([data_id["visit"] for data_id in d_nodes.keys()], [2] * 4)
138 # Finish the 'd' nodes and check that the walker iteration is done.
139 for n in d_nodes.values():
140 walker.finish(n)
141 with self.assertRaises(StopIteration):
142 next(walker)