Repository: PrefectHQ/prefect

Issue:
Update Docker healthcheck to include result handlers
There are certain situations that we can catch prior to deployment that would result in a Flow not behaving as expected.  In particular, if a user tries to deploy a Flow which:
- has a task with retry settings
- that task also has a keyed upstream edge
- the upstream task has no result handler

Then retries will fail to work in Cloud.  Because this situation can be inferred apriori we should raise an informative error during the Docker health check.

Similarly, if a user has tasks with `checkpoint=True` that don't have corresponding result handlers, we should raise an error.

Note that if the Flow itself has a result handler (but the Task does not), then all of this is OK because the Flow's should be the fallback.


Discussion:

Resolving patch:
diff --git a/src/prefect/environments/storage/_healthcheck.py b/src/prefect/environments/storage/_healthcheck.py
new file mode 100644
--- /dev/null
+++ b/src/prefect/environments/storage/_healthcheck.py
@@ -0,0 +1,102 @@
+"""
+A health check script which is intended to be executed inside a user's Docker container
+containing their Flow(s).  This class runs through a series of checks to catch possible
+deployment issues early on.
+"""
+
+import ast
+import cloudpickle
+import sys
+import warnings
+
+
+def system_check(python_version: str):
+    python_version = ast.literal_eval(python_version)  # convert string to tuple of ints
+
+    if (
+        sys.version_info.minor < python_version[1]
+        or sys.version_info.minor > python_version[1]
+    ):
+        msg = "Your Docker container is using python version {sys_ver}, but your Flow was serialized using {user_ver}; this could lead to unexpected errors in deployment.".format(
+            sys_ver=(sys.version_info.major, sys.version_info.minor),
+            user_ver=python_version,
+        )
+        warnings.warn(msg)
+    else:
+        print("System Version check: OK")
+
+
+def serialization_check(flow_file_paths: str):
+    flow_file_paths = ast.literal_eval(
+        flow_file_paths
+    )  # convert string to list of strings
+
+    flows = []
+    for flow_file in flow_file_paths:
+        with open(flow_file, "rb") as f:
+            flows.append(cloudpickle.load(f))
+
+    print("Serialization check: OK")
+    return flows
+
+
+def result_handler_check(flows: list):
+    for flow in flows:
+        if flow.result_handler is not None:
+            continue
+
+        ## test for tasks which are checkpointed with no result handler
+        if any([(t.checkpoint and t.result_handler is None) for t in flow.tasks]):
+            raise ValueError(
+                "Some tasks request to be checkpointed but do not have a result handler. See https://docs.prefect.io/guide/core_concepts/results.html for more details."
+            )
+
+        ## test for tasks which might retry without upstream result handlers
+        retry_tasks = [t for t in flow.tasks if t.max_retries > 0]
+        upstream_edges = flow.all_upstream_edges()
+        for task in retry_tasks:
+            if any(
+                [
+                    e.upstream_task.result_handler is None
+                    for e in upstream_edges[task]
+                    if e.key is not None
+                ]
+            ):
+                raise ValueError(
+                    "Task {} has retry settings but some upstream dependencies do not have result handlers. See https://docs.prefect.io/guide/core_concepts/results.html for more details.".format(
+                        task
+                    )
+                )
+
+        ## test for tasks which request caching with no result handler or no upstream result handlers
+        cached_tasks = [t for t in flow.tasks if t.cache_for is not None]
+        for task in cached_tasks:
+            if task.result_handler is None:
+                raise ValueError(
+                    "Task {} has cache settings but does not have a result handler. See https://docs.prefect.io/guide/core_concepts/results.html for more details.".format(
+                        task
+                    )
+                )
+            if any(
+                [
+                    e.upstream_task.result_handler is None
+                    for e in upstream_edges[task]
+                    if e.key is not None
+                ]
+            ):
+                raise ValueError(
+                    "Task {} has cache settings but some upstream dependencies do not have result handlers. See https://docs.prefect.io/guide/core_concepts/results.html for more details.".format(
+                        task
+                    )
+                )
+    print("Result Handler check: OK")
+
+
+if __name__ == "__main__":
+    flow_file_path, python_version = sys.argv[1:3]
+
+    print("Beginning health checks...")
+    system_check(python_version)
+    flows = serialization_check(flow_file_path)
+    result_handler_check(flows)
+    print("All health checks passed.")
diff --git a/src/prefect/environments/storage/docker.py b/src/prefect/environments/storage/docker.py
--- a/src/prefect/environments/storage/docker.py
+++ b/src/prefect/environments/storage/docker.py
@@ -251,7 +251,7 @@ def _build_image(self, push: bool = True) -> tuple:
 
             if len(client.images(name=full_name)) == 0:
                 raise SerializationError(
-                    "Your flow failed to deserialize in the container; please ensure that all necessary files and dependencies have been included."
+                    "Your flow failed one of its deployment health checks! Please ensure that all necessary files and dependencies have been included."
                 )
 
             # Push the image if requested
@@ -333,29 +333,10 @@ def create_dockerfile_object(self, directory: str = None) -> None:
                 )
 
             # Write a healthcheck script into the image
-            healthcheck = textwrap.dedent(
-                """\
-            print('Beginning health check...')
-            import cloudpickle
-            import sys
-            import warnings
-
-            for flow_file in [{flow_file_paths}]:
-                with open(flow_file, 'rb') as f:
-                    flow = cloudpickle.load(f)
-
-            if sys.version_info.minor < {python_version}[1] or sys.version_info.minor > {python_version}[1]:
-                msg = "Your Docker container is using python version {{sys_ver}}, but your Flow was serialized using {{user_ver}}; this could lead to unexpected errors in deployment.".format(sys_ver=(sys.version_info.major, sys.version_info.minor), user_ver={python_version})
-                warnings.warn(msg)
-
-            print('Healthcheck: OK')
-            """.format(
-                    flow_file_paths=", ".join(
-                        ["'{}'".format(k) for k in self.flows.values()]
-                    ),
-                    python_version=(sys.version_info.major, sys.version_info.minor),
-                )
-            )
+            with open(
+                os.path.join(os.path.dirname(__file__), "_healthcheck.py"), "r"
+            ) as healthscript:
+                healthcheck = healthscript.read()
 
             with open(os.path.join(directory, "healthcheck.py"), "w") as health_file:
                 health_file.write(healthcheck)
@@ -379,7 +360,7 @@ def create_dockerfile_object(self, directory: str = None) -> None:
                 RUN pip install git+https://github.com/PrefectHQ/prefect.git@{version}#egg=prefect[kubernetes]
                 # RUN pip install prefect
 
-                RUN python /root/.prefect/healthcheck.py
+                RUN python /root/.prefect/healthcheck.py '[{flow_file_paths}]' '{python_version}'
                 """.format(
                     base_image=self.base_image,
                     pip_installs=pip_installs,
@@ -387,6 +368,10 @@ def create_dockerfile_object(self, directory: str = None) -> None:
                     copy_files=copy_files,
                     env_vars=env_vars,
                     version=self.prefect_version,
+                    flow_file_paths=", ".join(
+                        ['"{}"'.format(k) for k in self.flows.values()]
+                    ),
+                    python_version=(sys.version_info.major, sys.version_info.minor),
                 )
             )
 


Test patch:
diff --git a/tests/environments/storage/test_defaults.py b/tests/environments/storage/test_defaults.py
--- a/tests/environments/storage/test_defaults.py
+++ b/tests/environments/storage/test_defaults.py
@@ -1,3 +1,5 @@
+import pytest
+
 from prefect.environments import storage
 from prefect import utilities
 
@@ -17,4 +19,6 @@ def test_default_storage_ignores_bad_config():
     with utilities.configuration.set_temporary_config(
         {"flows.defaults.storage.default_class": "FOOBAR"}
     ):
-        assert storage.get_default_storage_class() is storage.Docker
+
+        with pytest.warns(UserWarning):
+            assert storage.get_default_storage_class() is storage.Docker
diff --git a/tests/environments/storage/test_docker_healthcheck.py b/tests/environments/storage/test_docker_healthcheck.py
new file mode 100644
--- /dev/null
+++ b/tests/environments/storage/test_docker_healthcheck.py
@@ -0,0 +1,151 @@
+import cloudpickle
+import datetime
+import os
+import pytest
+import sys
+import tempfile
+
+from prefect import task, Flow, Task
+from prefect.environments.storage import _healthcheck as healthchecks
+
+
+class TestSerialization:
+    def test_serialization_check_raises_on_bad_imports(self):
+        bad_bytes = b"\x80\x04\x95\x18\x00\x00\x00\x00\x00\x00\x00\x8c\x0bfoo_machine\x94\x8c\x04func\x94\x93\x94."
+        with tempfile.NamedTemporaryFile() as f:
+            f.write(bad_bytes)
+            f.seek(0)
+            with pytest.raises(ImportError, match="foo_machine"):
+                objs = healthchecks.serialization_check("['{}']".format(f.name))
+
+    def test_serialization_check_passes_and_returns_objs(self):
+        good_bytes = cloudpickle.dumps(Flow("empty"))
+        with tempfile.NamedTemporaryFile() as f:
+            f.write(good_bytes)
+            f.seek(0)
+            objs = healthchecks.serialization_check("['{}']".format(f.name))
+
+        assert len(objs) == 1
+
+        flow = objs.pop()
+        assert isinstance(flow, Flow)
+        assert flow.name == "empty"
+        assert flow.tasks == set()
+
+    def test_serialization_check_passes_and_returns_multiple_objs(self):
+        flow_one = cloudpickle.dumps(Flow("one"))
+        flow_two = cloudpickle.dumps(Flow("two"))
+        with tempfile.TemporaryDirectory() as tmpdir:
+            file_one = os.path.join(tmpdir, "one.flow")
+            with open(file_one, "wb") as f:
+                f.write(flow_one)
+
+            file_two = os.path.join(tmpdir, "two.flow")
+            with open(file_two, "wb") as f:
+                f.write(flow_two)
+
+            paths = "['{0}', '{1}']".format(file_one, file_two)
+            objs = healthchecks.serialization_check(paths)
+
+        assert len(objs) == 2
+
+
+class TestSystemCheck:
+    def test_system_check_just_warns(self):
+        with pytest.warns(UserWarning, match="unexpected errors"):
+            healthchecks.system_check("(3, 4)")
+
+    def test_system_check_doesnt_warn(self):
+        sys_info = "({0}, {1})".format(sys.version_info.major, sys.version_info.minor)
+        with pytest.warns(None) as records:
+            healthchecks.system_check(sys_info)
+        assert len(records) == 0
+
+
+class TestResultHandlerCheck:
+    def test_no_raise_on_normal_flow(self):
+        @task
+        def up():
+            pass
+
+        @task
+        def down(x):
+            pass
+
+        with Flow("THIS IS A TEST") as flow:
+            result = down(x=up, upstream_tasks=[Task(), Task()])
+
+        assert healthchecks.result_handler_check([flow]) is None
+
+    @pytest.mark.parametrize(
+        "kwargs", [dict(checkpoint=True), dict(cache_for=datetime.timedelta(minutes=1))]
+    )
+    def test_raises_for_checkpointed_tasks(self, kwargs):
+        @task(**kwargs)
+        def up():
+            pass
+
+        f = Flow("foo-test", tasks=[up])
+
+        with pytest.raises(ValueError, match="have a result handler."):
+            healthchecks.result_handler_check([f])
+
+    @pytest.mark.parametrize(
+        "kwargs", [dict(checkpoint=True), dict(cache_for=datetime.timedelta(minutes=1))]
+    )
+    def test_doesnt_raise_for_checkpointed_tasks_if_flow_has_result_handler(
+        self, kwargs
+    ):
+        @task(**kwargs)
+        def up():
+            pass
+
+        f = Flow("foo-test", tasks=[up], result_handler=42)
+        assert healthchecks.result_handler_check([f]) is None
+
+    @pytest.mark.parametrize(
+        "kwargs",
+        [
+            dict(max_retries=2, retry_delay=datetime.timedelta(minutes=1)),
+            dict(cache_for=datetime.timedelta(minutes=1)),
+        ],
+    )
+    def test_raises_for_tasks_with_upstream_dependencies_with_no_result_handlers(
+        self, kwargs
+    ):
+        @task
+        def up():
+            pass
+
+        @task(**kwargs, result_handler=42)
+        def down(x):
+            pass
+
+        with Flow("upstream-test") as f:
+            result = down(x=up)
+
+        with pytest.raises(
+            ValueError, match="upstream dependencies do not have result handlers."
+        ):
+            healthchecks.result_handler_check([f])
+
+    @pytest.mark.parametrize(
+        "kwargs",
+        [
+            dict(max_retries=2, retry_delay=datetime.timedelta(minutes=1)),
+            dict(cache_for=datetime.timedelta(minutes=1)),
+        ],
+    )
+    def test_doesnt_raise_for_tasks_with_non_keyed_edges(self, kwargs):
+        @task
+        def up():
+            pass
+
+        @task(**kwargs, result_handler=42)
+        def down():
+            pass
+
+        with Flow("non-keyed-test") as f:
+            result = down(upstream_tasks=[up])
+
+        assert healthchecks.result_handler_check([f]) is None