Add a --viz_direction flag to CLI (nv-morpheus#1119)

* Add a `--viz_direction` flag to CLI
* Cherry-pick'd [9532bfe](nv-morpheus@9532bfe) from @mdemoret-nv 

fixes nv-morpheus#643

Authors:
  - David Gardner (https://github.com/dagardner-nv)
  - Michael Demoret (https://github.com/mdemoret-nv)

Approvers:
  - Christopher Harris (https://github.com/cwharris)
  - Michael Demoret (https://github.com/mdemoret-nv)

URL: nv-morpheus#1119

morpheus/cli/commands.py

@@ -53,6 +53,9 @@
 # autocomplete too much.
 FILE_TYPE_NAMES = ["auto", "csv", "json"]
 
+# Graphviz rankdir options ad documented in https://graphviz.org/docs/attr-types/rankdir/
+RANKDIR_CHOICES = ['BT', 'LR', 'RL', 'TB']
+
 ALIASES = {
     "pipeline": "pipeline-nlp",
 }
@@ -317,6 +320,11 @@ def run(ctx: click.Context, **kwargs):
               default=None,
               type=click.Path(dir_okay=False, writable=True),
               help="Save a visualization of the pipeline at the specified location")
+@click.option('--viz_direction',
+              default="LR",
+              type=click.Choice(RANKDIR_CHOICES, case_sensitive=False),
+              help=("Set the direction for the Graphviz pipeline diagram, "
+                    "ignored unless --viz_file is also specified."))
 @prepare_command()
 def pipeline_nlp(ctx: click.Context, **kwargs):
     """
@@ -382,6 +390,11 @@ def pipeline_nlp(ctx: click.Context, **kwargs):
               default=None,
               type=click.Path(dir_okay=False, writable=True),
               help="Save a visualization of the pipeline at the specified location")
+@click.option('--viz_direction',
+              default="LR",
+              type=click.Choice(RANKDIR_CHOICES, case_sensitive=False),
+              help=("Set the direction for the Graphviz pipeline diagram, "
+                    "ignored unless --viz_file is also specified."))
 @prepare_command()
 def pipeline_fil(ctx: click.Context, **kwargs):
     """
@@ -464,6 +477,11 @@ def pipeline_fil(ctx: click.Context, **kwargs):
               default=None,
               type=click.Path(dir_okay=False, writable=True),
               help="Save a visualization of the pipeline at the specified location")
+@click.option('--viz_direction',
+              default="LR",
+              type=click.Choice(RANKDIR_CHOICES, case_sensitive=False),
+              help=("Set the direction for the Graphviz pipeline diagram, "
+                    "ignored unless --viz_file is also specified."))
 @prepare_command()
 def pipeline_ae(ctx: click.Context, **kwargs):
     """
@@ -540,6 +558,11 @@ def pipeline_ae(ctx: click.Context, **kwargs):
               default=None,
               type=click.Path(dir_okay=False, writable=True),
               help="Save a visualization of the pipeline at the specified location")
+@click.option('--viz_direction',
+              default="LR",
+              type=click.Choice(RANKDIR_CHOICES, case_sensitive=False),
+              help=("Set the direction for the Graphviz pipeline diagram, "
+                    "ignored unless --viz_file is also specified."))
 @prepare_command()
 def pipeline_other(ctx: click.Context, **kwargs):
     """
@@ -604,7 +627,7 @@ def post_pipeline(ctx: click.Context, *args, **kwargs):
 
     # TODO(MDD): Move visualization before `pipeline.run()` once Issue #230 is fixed.
     if ("viz_file" in kwargs and kwargs["viz_file"] is not None):
-        pipeline.visualize(kwargs["viz_file"], rankdir="LR")
+        pipeline.visualize(kwargs["viz_file"], rankdir=kwargs["viz_direction"].upper())
         click.secho(f"Pipeline visualization saved to {kwargs['viz_file']}", fg="yellow")
 
 

tests/test_pipe_viz.py

@@ -16,23 +16,27 @@
 
 import imghdr
 import os
+import subprocess
 
 import pytest
 
+from morpheus.cli.commands import RANKDIR_CHOICES
 from morpheus.pipeline import LinearPipeline
 from morpheus.pipeline.pipeline import Pipeline
 from morpheus.stages.input.in_memory_source_stage import InMemorySourceStage
 from morpheus.stages.output.in_memory_sink_stage import InMemorySinkStage
 from morpheus.stages.postprocess.add_classifications_stage import AddClassificationsStage
 from morpheus.stages.postprocess.serialize_stage import SerializeStage
 from morpheus.stages.preprocess.deserialize_stage import DeserializeStage
+from utils import TEST_DIRS
 from utils import assert_path_exists
+from utils.dataset_manager import DatasetManager
 from utils.stages.conv_msg import ConvMsg
 
 
 @pytest.mark.use_cudf
-@pytest.fixture(scope="function")
-def viz_pipeline(config, filter_probs_df):
+@pytest.fixture(name="viz_pipeline", scope="function")
+def viz_pipeline_fixture(config, filter_probs_df):
     """
     Creates a quick pipeline.
     """
@@ -44,13 +48,13 @@ def viz_pipeline(config, filter_probs_df):
     pipe.add_stage(DeserializeStage(config))
     pipe.add_stage(ConvMsg(config, filter_probs_df))
     pipe.add_stage(AddClassificationsStage(config))
-    pipe.add_stage(SerializeStage(config, include=["^{}$".format(c) for c in config.class_labels]))
+    pipe.add_stage(SerializeStage(config, include=[f"^{c}$" for c in config.class_labels]))
     pipe.add_stage(InMemorySinkStage(config))
 
     return pipe
 
 
-def test_call_before_run(viz_pipeline: Pipeline, tmp_path):
+def test_call_before_run(viz_pipeline: Pipeline, tmp_path: str):
 
     # Test is necessary to ensure run() is called first. See issue #230
     viz_file = os.path.join(tmp_path, 'pipeline.png')
@@ -60,7 +64,7 @@ def test_call_before_run(viz_pipeline: Pipeline, tmp_path):
         viz_pipeline.visualize(viz_file, rankdir="LR")
 
 
-def test_png(viz_pipeline: Pipeline, tmp_path):
+def test_png(viz_pipeline: Pipeline, tmp_path: str):
 
     viz_file = os.path.join(tmp_path, 'pipeline.png')
 
@@ -72,3 +76,21 @@ def test_png(viz_pipeline: Pipeline, tmp_path):
     # Verify that the output file exists and is a valid png file
     assert_path_exists(viz_file)
     assert imghdr.what(viz_file) == 'png'
+
+
+@pytest.mark.slow
+@pytest.mark.parametrize("rankdir", RANKDIR_CHOICES)
+def test_from_cli(tmp_path: str, dataset_pandas: DatasetManager, rankdir: str):
+    input_file = os.path.join(TEST_DIRS.tests_data_dir, 'filter_probs.csv')
+    out_file = os.path.join(tmp_path, 'out.csv')
+    viz_file = os.path.join(tmp_path, 'pipeline.png')
+    cli = (f"morpheus run pipeline-other --viz_file={viz_file} --viz_direction={rankdir} "
+           f"from-file --filename={input_file} to-file --filename={out_file}")
+    subprocess.run(cli, check=True, shell=True)
+
+    assert_path_exists(viz_file)
+    assert imghdr.what(viz_file) == 'png'
+    assert_path_exists(out_file)
+
+    df = dataset_pandas.get_df(out_file, no_cache=True)
+    dataset_pandas.assert_compare_df(df, dataset_pandas['filter_probs.csv'])