Merge branch 'tickets/DM-44647'

lsst · Jun 26, 2024 · 2883a1f · 2883a1f
2 parents f1e75d8 + 28d8a12
commit 2883a1f
Show file tree

Hide file tree

Showing 6 changed files with 31 additions and 12 deletions.
diff --git a/doc/changes/DM-44647.feature.md b/doc/changes/DM-44647.feature.md
@@ -0,0 +1,6 @@
+Move pipeline-dot build from cmdLineFwk to builder.
+
+This PR moves the pipeline-dot build from the cmdLineFwk package to the builder package.
+This is done to make the pipeline-dot build more accessible to other packages.
+As part of this change, output pipeline-dot files contain dimensions and storage classes for each dataset.
+This change also includes updates to existing unit tests to reflect the new output types.
diff --git a/python/lsst/ctrl/mpexec/cli/cmd/commands.py b/python/lsst/ctrl/mpexec/cli/cmd/commands.py
@@ -124,10 +124,12 @@ def build(ctx: click.Context, **kwargs: Any) -> None:
     """
     kwargs = _collectActions(ctx, **kwargs)
     show = ShowInfo(kwargs.pop("show", []))
-    if kwargs.get("butler_config") is not None and {"pipeline-graph", "task-graph"}.isdisjoint(show.commands):
+    if kwargs.get("butler_config") is not None and (
+        {"pipeline-graph", "task-graph"}.isdisjoint(show.commands) and not kwargs.get("pipeline_dot")
+    ):
         raise click.ClickException(
             "--butler-config was provided but nothing uses it "
-            "(only --show pipeline-graph and --show task-graph do)."
+            "(only --show pipeline-graph, --show task-graph and --pipeline-dot do)."
         )
     script.build(**kwargs, show=show)
     _unhandledShow(show, "build")

diff --git a/python/lsst/ctrl/mpexec/cli/script/build.py b/python/lsst/ctrl/mpexec/cli/script/build.py
@@ -28,6 +28,7 @@
 from types import SimpleNamespace
 
 from lsst.daf.butler import Butler
+from lsst.pipe.base.pipeline_graph import visualization
 
 from ... import CmdLineFwk
 from ..utils import _PipelineAction
@@ -114,6 +115,15 @@ def build(  # type: ignore
     else:
         butler = None
 
+    if pipeline_dot:
+        with open(pipeline_dot, "w") as stream:
+            visualization.show_dot(
+                pipeline.to_graph(butler.registry if butler is not None else None, visualization_only=True),
+                stream,
+                dataset_types=True,
+                task_classes="full",
+            )
+
     show.show_pipeline_info(pipeline, butler=butler)
 
     return pipeline
diff --git a/python/lsst/ctrl/mpexec/cmdLineFwk.py b/python/lsst/ctrl/mpexec/cmdLineFwk.py
@@ -75,7 +75,7 @@
 from lsst.utils.logging import getLogger
 from lsst.utils.threads import disable_implicit_threading
 
-from .dotTools import graph2dot, pipeline2dot
+from .dotTools import graph2dot
 from .executionGraphFixup import ExecutionGraphFixup
 from .mpGraphExecutor import MPGraphExecutor
 from .preExecInit import PreExecInit, PreExecInitLimited
@@ -582,9 +582,6 @@ def makePipeline(self, args: SimpleNamespace) -> Pipeline:
         if args.save_pipeline:
             pipeline.write_to_uri(args.save_pipeline)
 
-        if args.pipeline_dot:
-            pipeline2dot(pipeline, args.pipeline_dot)
-
         return pipeline
 
     def makeGraph(self, pipeline: Pipeline, args: SimpleNamespace) -> QuantumGraph | None:

diff --git a/python/lsst/ctrl/mpexec/showInfo.py b/python/lsst/ctrl/mpexec/showInfo.py
@@ -172,9 +172,13 @@ def show_pipeline_info(self, pipeline: Pipeline, butler: Butler | None) -> None:
                 case "tasks":
                     self._showTaskHierarchy(pipeline)
                 case "pipeline-graph":
-                    visualization.show(pipeline.to_graph(registry), self.stream, dataset_types=True)
+                    visualization.show(
+                        pipeline.to_graph(registry, visualization_only=True), self.stream, dataset_types=True
+                    )
                 case "task-graph":
-                    visualization.show(pipeline.to_graph(registry), self.stream, dataset_types=False)
+                    visualization.show(
+                        pipeline.to_graph(registry, visualization_only=True), self.stream, dataset_types=False
+                    )
                 case _:
                     raise RuntimeError(f"Unexpectedly tried to process command {command!r}.")
             self.handled.add(command)

diff --git a/tests/test_cmdLineFwk.py b/tests/test_cmdLineFwk.py
@@ -478,11 +478,11 @@ def testShowPipeline(self):
         self.assertEqual(
             "\n".join(
                 [
-                    "○  add_dataset_in",
+                    "○  add_dataset_in: {detector} NumpyArray",
                     "│",
-                    "■  task",
+                    "■  task: {detector}",
                     "│",
-                    "◍  add_dataset_out, add2_dataset_out",
+                    "◍  add_dataset_out, add2_dataset_out: {detector} NumpyArray",
                 ]
             ),
             output,
@@ -493,7 +493,7 @@ def testShowPipeline(self):
         show.show_pipeline_info(pipeline, None)
         stream.seek(0)
         output = stream.read().strip()
-        self.assertEqual("■  task", output)
+        self.assertEqual("■  task: {detector}", output)
 
         stream = StringIO()
         show = ShowInfo(["config=task::addEnd"], stream=stream)  # Match but warns