Repository: PrefectHQ/prefect

Issue:
Saving the output of `flow.visualize`
Prefect currently generates a random filename in order to save the state of the flow. It's more useful to store the visualization to a know location so that we can consume the file downstream (either upload it to S3 or push it to a monitoring service, etc.)


Discussion:

Resolving patch:
diff --git a/examples/flow_state_handler_viz.py b/examples/flow_state_handler_viz.py
new file mode 100644
--- /dev/null
+++ b/examples/flow_state_handler_viz.py
@@ -0,0 +1,50 @@
+"""
+A simple example Flow which stores the current flow visualization to a file
+each time the flow changes state, using a Flow-level state handler.
+
+Uses the same flow from "Retries w/ Mapping" on a minute schedule.
+"""
+import random
+from datetime import datetime, timedelta
+
+from prefect import Flow, task
+from prefect.schedules import IntervalSchedule
+
+
+def visualize(flow, old_state, new_state):
+    """
+    Flow state handler which stores the current flow
+    visualization to a known file location, for live viewing.
+    """
+    if isinstance(new_state.result, dict):
+        ## note that the relevant file will actually be called "known_location.pdf"
+        flow.visualize(flow_state=new_state, filename="known_location")
+    return new_state
+
+
+@task
+def generate_random_list():
+    n = random.randint(15, 25)
+    return list(range(n))
+
+
+@task(max_retries=3, retry_delay=timedelta(seconds=0))
+def randomly_fail():
+    x = random.random()
+    if x > 0.7:
+        raise ValueError("x is too large")
+
+
+schedule = IntervalSchedule(
+    start_date=datetime.utcnow(),
+    interval=timedelta(minutes=1),
+    end_date=datetime.utcnow() + timedelta(minutes=10),
+)
+
+with Flow(
+    "random-mapping-with-viz", schedule=schedule, state_handlers=[visualize]
+) as f:
+    randomly_fail.map(upstream_tasks=[generate_random_list])
+
+
+f.run()
diff --git a/src/prefect/core/flow.py b/src/prefect/core/flow.py
--- a/src/prefect/core/flow.py
+++ b/src/prefect/core/flow.py
@@ -1000,14 +1000,19 @@ def run(
 
     # Visualization ------------------------------------------------------------
 
-    def visualize(self, flow_state: "prefect.engine.state.State" = None) -> object:
+    def visualize(
+        self, flow_state: "prefect.engine.state.State" = None, filename: str = None
+    ) -> object:
         """
         Creates graphviz object for representing the current flow; this graphviz
         object will be rendered inline if called from an IPython notebook, otherwise
-        it will be rendered in a new window.
+        it will be rendered in a new window.  If a `filename` is provided, the object
+        will not be rendered and instead saved to the location specified.
 
         Args:
             - flow_state (State, optional): flow state object used to optionally color the nodes
+            - filename (str, optional): a filename specifying a location to save this visualization to; if provided,
+                the visualization will not be rendered automatically
 
         Raises:
             - ImportError: if `graphviz` is not installed
@@ -1093,12 +1098,15 @@ def get_color(task: Task, map_index: int = None) -> str:
         except Exception:
             pass
 
-        with tempfile.NamedTemporaryFile(delete=False) as tmp:
-            tmp.close()
-            try:
-                graph.render(tmp.name, view=True)
-            finally:
-                os.unlink(tmp.name)
+        if filename:
+            graph.render(filename, view=False)
+        else:
+            with tempfile.NamedTemporaryFile(delete=False) as tmp:
+                tmp.close()
+                try:
+                    graph.render(tmp.name, view=True)
+                finally:
+                    os.unlink(tmp.name)
 
         return graph
 


Test patch:
diff --git a/tests/core/test_flow.py b/tests/core/test_flow.py
--- a/tests/core/test_flow.py
+++ b/tests/core/test_flow.py
@@ -963,6 +963,20 @@ def test_viz_returns_graph_object_if_in_ipython(self):
         assert "label=a_nice_task" in graph.source
         assert "shape=ellipse" in graph.source
 
+    def test_viz_saves_graph_object_if_filename(self):
+        import graphviz
+
+        f = Flow(name="test")
+        f.add_task(Task(name="a_nice_task"))
+
+        with tempfile.NamedTemporaryFile() as tmp:
+            graph = f.visualize(filename=tmp.name)
+            with open(tmp.name, "r") as f:
+                contents = f.read()
+
+        assert "label=a_nice_task" in contents
+        assert "shape=ellipse" in contents
+
     def test_viz_reflects_mapping(self):
         ipython = MagicMock(
             get_ipython=lambda: MagicMock(config=dict(IPKernelApp=True))
diff --git a/tests/utilities/test_serialization.py b/tests/utilities/test_serialization.py
--- a/tests/utilities/test_serialization.py
+++ b/tests/utilities/test_serialization.py
@@ -1,4 +1,5 @@
 import datetime
+import itertools
 import uuid
 
 import marshmallow
@@ -234,11 +235,14 @@ def test_serialize_outer_no_state(self):
         assert serialized["f"] == "tests.utilities.test_serialization.outer()"
 
     def test_serialize_outer_with_state(self):
+        """Have to account for order because of Python 3.5"""
         serialized = self.Schema().dump(dict(f=outer(x=1, y=2, z=99)))
-        assert (
-            serialized["f"]
-            == "tests.utilities.test_serialization.outer(x=1, y=2, z=99)"
-        )
+        endings = list(itertools.permutations(["x=1", "y=2", "z=99"]))
+        valid = [
+            "tests.utilities.test_serialization.outer(" + ", ".join(sig) + ")"
+            for sig in endings
+        ]
+        assert serialized["f"] in valid
 
     def test_serialize_invalid_fn(self):
         with pytest.raises(marshmallow.ValidationError):