Coverage for python/lsst/ctrl/bps/clustered_quantum

1# This file is part of ctrl_bps.

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.

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"""Class definitions for a Clustered QuantumGraph where a node in the graph is

29a QuantumGraph.

30"""

32from __future__ import annotations

34__all__ = ["ClusteredQuantumGraph", "QuantaCluster"]

36import logging

37import pickle

38import re

39import uuid

40from collections import Counter, defaultdict

41from pathlib import Path

43from networkx import DiGraph, is_directed_acyclic_graph, is_isomorphic, topological_sort

45from lsst.pipe.base.pipeline_graph import TaskImportMode

46from lsst.pipe.base.quantum_graph import PredictedQuantumGraph, QuantumInfo

47from lsst.utils.iteration import ensure_iterable

49from .bps_draw import draw_networkx_dot

51_LOG = logging.getLogger(__name__)

54class QuantaCluster:

55 """Information about the cluster and Quanta belonging to it.

57 Parameters

58 ----------

59 name : `str`

60 Lookup key (logical file name) of file/directory. Must

61 be unique within ClusteredQuantumGraph.

62 label : `str`

63 Value used to group clusters.

64 tags : `dict` [`str`, `Any`], optional

65 Arbitrary key/value pairs for the cluster.

67 Raises

68 ------

69 ValueError

70 Raised if invalid name (e.g., name contains /).

71 """

73 def __init__(self, name, label, tags=None):

74 if "/" in name:

75 raise ValueError(f"Cluster's name cannot have a / ({name})")

76 self.name = name

77 self.label = label

78 self._qgraph_node_ids = []

79 self._task_label_counts = Counter()

80 self.tags = tags

81 if self.tags is None:

82 self.tags = {}

84 @classmethod

85 def from_quantum_info(

86 cls, quantum_id: uuid.UUID, quantum_info: QuantumInfo, template: str

87 ) -> QuantaCluster:

88 """Create single quantum cluster from the given quantum information.

90 Parameters

91 ----------

92 quantum_id : `uuid.UUID`

93 ID of the quantum.

94 quantum_info : `lsst.pipe.base.quantum_graph.QuantumInfo`

95 Dictionary of additional information about the quantum.

96 template : `str`

97 Template for creating cluster name.

99 Returns

100 -------

101 cluster : `QuantaCluster`

102 Newly created cluster containing the given quantum.

103 """

104 label = quantum_info["task_label"]

105 data_id = quantum_info["data_id"]

106

107 # Gather info for name template into a dictionary.

108 info = dict(data_id.required)

109 info["label"] = label

110 info["node_number"] = quantum_id

111 _LOG.debug("template = %s", template)

112 _LOG.debug("info for template = %s", info)

113

114 # Use dictionary plus template format string to create name. To avoid

115 # key errors from generic patterns, use defaultdict.

116 try:

117 name = template.format_map(defaultdict(lambda: "", info))

118 except TypeError:

119 _LOG.error("Problems creating cluster name. template='%s', info=%s", template, info)

120 raise

121 name = re.sub("_+", "_", name)

122 _LOG.debug("template name = %s", name)

123

124 cluster = QuantaCluster(name, label, info)

125 cluster.add_quantum(quantum_id, label)

126 return cluster

127

128 @property

129 def qgraph_node_ids(self):

130 """Quantum graph NodeIds corresponding to this cluster."""

131 _LOG.debug("_qgraph_node_ids = %s", self._qgraph_node_ids)

132 return frozenset(self._qgraph_node_ids)

133

134 @property

135 def quanta_counts(self):

136 """Counts of Quanta per taskDef.label in this cluster."""

137 return Counter(self._task_label_counts)

138

139 def add_quantum(self, node_id, task_label):

140 """Add a quantumNode to this cluster.

141

142 Parameters

143 ----------

144 node_id : `uuid.UUID`

145 ID for quantumNode to be added to cluster.

146 task_label : `str`

147 Task label for quantumNode to be added to cluster.

148 """

149 self._qgraph_node_ids.append(node_id)

150 self._task_label_counts[task_label] += 1

151

152 def __str__(self):

153 return (

154 f"QuantaCluster(name={self.name},label={self.label},tags={self.tags},"

155 f"counts={self.quanta_counts},ids={self.qgraph_node_ids})"

156 )

157

158 def __eq__(self, other: object) -> bool:

159 # Doesn't check data equality, but only

160 # name equality since those are supposed

161 # to be unique.

162 if isinstance(other, str):

163 return self.name == other

164

165 if isinstance(other, QuantaCluster):

166 return self.name == other.name

167

168 return False

169

170 def __hash__(self) -> int:

171 return hash(self.name)

172

173

174class ClusteredQuantumGraph:

175 """Graph where the data for a node is a subgraph of the full

176 QuantumGraph represented by a list of node ids.

177

178 Parameters

179 ----------

180 name : `str`

181 Name to be given to the ClusteredQuantumGraph.

182 qgraph : `lsst.pipe.base.quantum_graph.PredictedQuantumGraph`

183 The quantum graph to be clustered.

184 qgraph_filename : `str`

185 Filename for given quantum graph.

186

187 Raises

188 ------

189 ValueError

190 Raised if invalid name (e.g., name contains /)

191

192 Notes

193 -----

194 Using lsst.pipe.base.NodeId instead of integer because the QuantumGraph

195 API requires them. Chose skipping the repeated creation of objects to

196 use API over totally minimized memory usage.

197 """

198

199 def __init__(self, name: str, qgraph: PredictedQuantumGraph, qgraph_filename: str):

200 if "/" in name:

201 raise ValueError(f"name cannot have a / ({name})")

202 self._name = name

203 self._quantum_graph = qgraph

204 self._quantum_only_xgraph = qgraph.quantum_only_xgraph

205 self._quantum_graph_filename = Path(qgraph_filename).resolve()

206 self._cluster_graph = DiGraph()

207

208 def __str__(self):

209 return (

210 f"ClusteredQuantumGraph(name={self.name},"

211 f"quantum_graph_filename={self._quantum_graph_filename},"

212 f"len(qgraph)={len(self._quantum_graph) if self._quantum_graph else None},"

213 f"len(cqgraph)={len(self._cluster_graph) if self._cluster_graph else None})"

214 )

215

216 def __len__(self):

217 """Return the number of clusters."""

218 return len(self._cluster_graph)

219

220 def __eq__(self, other):

221 if not isinstance(other, ClusteredQuantumGraph):

222 return False

223 if len(self) != len(other):

224 return False

225 return is_isomorphic(self.qxgraph, other.qxgraph) and is_isomorphic(

226 self._cluster_graph, other._cluster_graph

227 )

228

229 @property

230 def name(self) -> str:

231 """The name of the ClusteredQuantumGraph."""

232 return self._name

233

234 @property

235 def qgraph(self) -> PredictedQuantumGraph:

236 """The quantum graph associated with this Clustered

237 QuantumGraph.

238 """

239 return self._quantum_graph

240

241 @property

242 def qxgraph(self) -> DiGraph:

243 """A networkx graph of all quanta."""

244 return self._quantum_only_xgraph

245

246 def add_cluster(self, clusters_for_adding):

247 """Add a cluster of quanta as a node in the graph.

248

249 Parameters

250 ----------

251 clusters_for_adding : `QuantaCluster` or `Iterable` [`QuantaCluster`]

252 The cluster to be added to the ClusteredQuantumGraph.

253 """

254 for cluster in ensure_iterable(clusters_for_adding):

255 if not isinstance(cluster, QuantaCluster):

256 raise TypeError(f"Must be type QuantaCluster (given: {type(cluster)})")

257

258 if self._cluster_graph.has_node(cluster.name):

259 raise KeyError(f"Cluster {cluster.name} already exists in ClusteredQuantumGraph")

260

261 self._cluster_graph.add_node(cluster.name, cluster=cluster)

262

263 def get_cluster(self, name):

264 """Retrieve a cluster from the ClusteredQuantumGraph by name.

265

266 Parameters

267 ----------

268 name : `str`

269 Name of cluster to retrieve.

270

271 Returns

272 -------

273 cluster : `QuantaCluster`

274 QuantaCluster matching given name.

275

276 Raises

277 ------

278 KeyError

279 Raised if the ClusteredQuantumGraph does not contain

280 a cluster with given name.

281 """

282 try:

283 attr = self._cluster_graph.nodes[name]

284 except KeyError as ex:

285 raise KeyError(f"{self.name} does not have a cluster named {name}") from ex

286 return attr["cluster"]

287

288 def get_quantum_info(self, id_: uuid.UUID) -> QuantumInfo:

289 """Retrieve a quantum info dict from the ClusteredQuantumGraph by ID.

290

291 Parameters

292 ----------

293 id_ : `uuid.UUID`

294 ID of the quantum to retrieve.

295

296 Returns

297 -------

298 quantum_info : `lsst.pipe.base.quantum_graph.QuantumInfo`

299 Quantum info dictionary for the given ID.

300

301 Raises

302 ------

303 KeyError

304 Raised if the ClusteredQuantumGraph does not contain

305 a quantum with given ID.

306 """

307 return self._quantum_only_xgraph.nodes[id_]

308

309 def __iter__(self):

310 """Iterate over names of clusters.

311

312 Returns

313 -------

314 names : `Iterator` [`str`]

315 Iterator over names of clusters.

316 """

317 return self._cluster_graph.nodes()

318

319 def clusters(self):

320 """Iterate over clusters.

321

322 Returns

323 -------

324 clusters : `Iterator` [`lsst.ctrl.bps.QuantaCluster`]

325 Iterator over clusters in topological order.

326 """

327 return map(self.get_cluster, topological_sort(self._cluster_graph))

328

329 def successors(self, name):

330 """Return clusters that are successors of the cluster

331 with the given name.

332

333 Parameters

334 ----------

335 name : `str`

336 Name of cluster for which need the successors.

337

338 Returns

339 -------

340 clusters : `Iterator` [`lsst.ctrl.bps.QuantaCluster`]

341 Iterator over successors of given cluster.

342 """

343 return map(self.get_cluster, self._cluster_graph.successors(name))

344

345 def predecessors(self, name):

346 """Return clusters that are predecessors of the cluster

347 with the given name.

348

349 Parameters

350 ----------

351 name : `str`

352 Name of cluster for which need the predecessors.

353

354 Returns

355 -------

356 clusters : `Iterator` [`lsst.ctrl.bps.QuantaCluster`]

357 Iterator over predecessors of given cluster.

358 """

359 return map(self.get_cluster, self._cluster_graph.predecessors(name))

360

361 def add_dependency(self, parent, child):

362 """Add a directed dependency between a parent cluster and a child

363 cluster.

364

365 Parameters

366 ----------

367 parent : `str` or `QuantaCluster`

368 Parent cluster.

369 child : `str` or `QuantaCluster`

370 Child cluster.

371

372 Raises

373 ------

374 KeyError

375 Raised if either the parent or child doesn't exist in the

376 ClusteredQuantumGraph.

377 """

378 if not self._cluster_graph.has_node(parent):

379 raise KeyError(f"{self.name} does not have a cluster named {parent}")

380 if not self._cluster_graph.has_node(child):

381 raise KeyError(f"{self.name} does not have a cluster named {child}")

382 _LOG.debug("add_dependency: adding edge %s %s", parent, child)

383

384 if isinstance(parent, QuantaCluster):

385 pname = parent.name

386 else:

387 pname = parent

388

389 if isinstance(child, QuantaCluster):

390 cname = child.name

391 else:

392 cname = child

393 self._cluster_graph.add_edge(pname, cname)

394

395 def __contains__(self, name):

396 """Check if a cluster with given name is in this ClusteredQuantumGraph.

397

398 Parameters

399 ----------

400 name : `str`

401 Name of cluster to check.

402

403 Returns

404 -------

405 found : `bool`

406 Whether a cluster with given name is in this ClusteredQuantumGraph.

407 """

408 return self._cluster_graph.has_node(name)

409

410 def save(self, filename, format_=None):

411 """Save the ClusteredQuantumGraph in a format that is loadable.

412

413 The quantum graph is assumed to have been saved separately.

414

415 Parameters

416 ----------

417 filename : `str`

418 File to which the ClusteredQuantumGraph should be serialized.

419 format_ : `str`, optional

420 Format in which to write the data. It defaults to pickle format.

421 """

422 path = Path(filename)

423

424 # if format is None, try extension

425 if format_ is None:

426 format_ = path.suffix[1:] # suffix includes the leading period

427

428 if format_ not in {"pickle"}:

429 raise RuntimeError(f"Unknown format ({format_})")

430

431 if format_ == "pickle":

432 # Don't save QuantumGraph in same file.

433 tmp_qgraph = self._quantum_graph

434 self._quantum_graph = None

435 with open(filename, "wb") as fh:

436 pickle.dump(self, fh)

437 # Return to original state.

438 self._quantum_graph = tmp_qgraph

439

440 def draw(self, filename, format_=None):

441 """Draw the ClusteredQuantumGraph in a given format.

442

443 Parameters

444 ----------

445 filename : `str`

446 File to which the ClusteredQuantumGraph should be serialized.

447 format_ : `str`, optional

448 Format in which to draw the data. It defaults to dot format.

449 """

450 path = Path(filename)

451

452 # if format is None, try extension

453 if format_ is None:

454 format_ = path.suffix[1:] # suffix includes the leading period

455

456 draw_funcs = {"dot": draw_networkx_dot}

457 if format_ in draw_funcs:

458 draw_funcs[format_](self._cluster_graph, filename)

459 else:

460 raise RuntimeError(f"Unknown draw format ({format_})")

461

462 @classmethod

463 def load(cls, filename, format_=None):

464 """Load a ClusteredQuantumGraph from the given file.

465

466 Parameters

467 ----------

468 filename : `str`

469 File from which to read the ClusteredQuantumGraph.

470 format_ : `str`, optional

471 Format of data to expect when loading from stream. It defaults

472 to pickle format.

473

474 Returns

475 -------

476 ClusteredQuantumGraph : `lsst.ctrl.bps.ClusteredQuantumGraph`

477 ClusteredQuantumGraph workflow loaded from the given file.

478 The QuantumGraph is loaded from its own file specified in

479 the saved ClusteredQuantumGraph.

480 """

481 path = Path(filename)

482

483 # if format is None, try extension

484 if format_ is None:

485 format_ = path.suffix[1:] # suffix includes the leading period

486

487 if format_ not in {"pickle"}:

488 raise RuntimeError(f"Unknown format ({format_})")

489

490 cgraph = None

491 if format_ == "pickle":

492 with open(filename, "rb") as fh:

493 cgraph: ClusteredQuantumGraph = pickle.load(fh)

494

495 # The QuantumGraph was saved separately

496 with PredictedQuantumGraph.open(

497 cgraph._quantum_graph_filename, import_mode=TaskImportMode.DO_NOT_IMPORT

498 ) as reader:

499 reader.read_thin_graph()

500 cgraph._quantum_graph = reader.finish()

501

502 return cgraph

503

504 def validate(self):

505 """Check correctness of completed ClusteredQuantumGraph.

506

507 Raises

508 ------

509 RuntimeError

510 If the ClusteredQuantumGraph is not valid.

511 """

512 # Check no cycles

513 if not is_directed_acyclic_graph(self._cluster_graph):

514 raise RuntimeError("ClusteredQuantumGraph is not a directed acyclic graph.")

515

516 # Check that Quantum only in 1 cluster

517 # Check cluster tags label matches cluster label

518 node_ids = set()

519 for cluster in self.clusters():

520 if "label" in cluster.tags and cluster.tags["label"] != cluster.label:

521 raise RuntimeError(

522 f"Label mismatch in cluster {cluster.name}: "

523 f"cluster={cluster.label} tags={cluster.tags['label']}"

524 )

525

526 for node_id in cluster.qgraph_node_ids:

527 if node_id in node_ids:

528 raise RuntimeError(

529 f"Quantum {node_id} occurs in at least 2 clusters (one of which is {cluster.name})"

530 )

531 else:

532 node_ids.add(node_id)

533

534 # Check that have all Quanta

535 quanta_count_qgraph = len(self._quantum_graph)

536 quanta_count_cqgraph = len(node_ids)

537 if quanta_count_qgraph != quanta_count_cqgraph:

538 raise RuntimeError(

539 f"Number of Quanta in clustered qgraph ({quanta_count_cqgraph}) does not equal number in"

540 f" quantum graph ({quanta_count_qgraph})"

541 )

Coverage for python / lsst / ctrl / bps / clustered_quantum_graph.py: 22%

180 statements