Repository: PrefectHQ/prefect

Issue:
Serialize State in New Triggers
With #947 merged, we now have triggers (and have always had `cache_validators`) which carry _state_.  However, when we serialize them to Cloud, we only carry their names for display in the UI.

It would be a very nice-to-have if we could serialize the state as well, i.e., serialize the bounds in the triggers and the validation lists in the curried cache validators for display in the UI.


Discussion:

Resolving patch:
diff --git a/src/prefect/engine/cache_validators.py b/src/prefect/engine/cache_validators.py
--- a/src/prefect/engine/cache_validators.py
+++ b/src/prefect/engine/cache_validators.py
@@ -13,10 +13,9 @@
 
 A cache validator returns `True` if the cache is still valid, and `False` otherwise.
 """
-from typing import Any, Dict, Iterable
+from typing import Any, Callable, Dict, Iterable
 
 import pendulum
-from toolz import curry
 
 import prefect
 
@@ -122,28 +121,17 @@ def all_parameters(
         return False
 
 
-@curry
-def partial_parameters_only(
-    state: "prefect.engine.state.Cached",
-    inputs: Dict[str, Any],
-    parameters: Dict[str, Any],
-    validate_on: Iterable[str] = None,
-) -> bool:
+def partial_parameters_only(validate_on: Iterable[str] = None,) -> Callable:
     """
     Validates the cache based on cache expiration _and_ a subset of parameters (determined by the
     `validate_on` keyword) which were provided on the last successful run.
 
     Args:
-        - state (State): a `Success` state from the last successful Task run which contains the cache
-        - inputs (dict): a `dict` of inputs which were available on the last
-            successful run of the cached Task
-        - parameters (dict): a `dict` of parameters which were available on the
-            last successful run of the cached Task
         - validate_on (list): a `list` of strings specifying the parameter names
             to validate against
 
     Returns:
-        - boolean specifying whether or not the cache should be used
+        - Callable: the actual validation function specifying whether or not the cache should be used
 
     Example:
     ```python
@@ -170,44 +158,56 @@ def daily_db_refresh(nrows, runtime):
                    task_states={result: state1.result[db_state]})
     ```
     """
-    parameters = parameters or {}
-    if duration_only(state, inputs, parameters) is False:
-        return False
-    elif validate_on is None:
-        return True  # if you dont want to validate on anything, then the cache is valid
-    else:
-        cached = state.cached_parameters or {}
-        partial_provided = {
-            key: value for key, value in parameters.items() if key in validate_on
-        }
-        partial_needed = {
-            key: value for key, value in cached.items() if key in validate_on
-        }
-        return partial_provided == partial_needed
-
-
-@curry
-def partial_inputs_only(
-    state: "prefect.engine.state.Cached",
-    inputs: Dict[str, Any],
-    parameters: Dict[str, Any],
-    validate_on: Iterable[str] = None,
-) -> bool:
+
+    def _partial_parameters_only(
+        state: "prefect.engine.state.Cached",
+        inputs: Dict[str, Any],
+        parameters: Dict[str, Any],
+    ) -> bool:
+        """
+        The actual cache validation function which will be used.
+
+        Args:
+            - state (State): a `Success` state from the last successful Task run which contains the cache
+            - inputs (dict): a `dict` of inputs which were available on the last
+                successful run of the cached Task
+            - parameters (dict): a `dict` of parameters which were available on the
+                last successful run of the cached Task
+
+        Returns:
+            - boolean specifying whether or not the cache should be used
+        """
+        parameters = parameters or {}
+        if duration_only(state, inputs, parameters) is False:
+            return False
+        elif validate_on is None:
+            return (
+                True
+            )  # if you dont want to validate on anything, then the cache is valid
+        else:
+            cached = state.cached_parameters or {}
+            partial_provided = {
+                key: value for key, value in parameters.items() if key in validate_on
+            }
+            partial_needed = {
+                key: value for key, value in cached.items() if key in validate_on
+            }
+            return partial_provided == partial_needed
+
+    return _partial_parameters_only
+
+
+def partial_inputs_only(validate_on: Iterable[str] = None,) -> Callable:
     """
     Validates the cache based on cache expiration _and_ a subset of inputs (determined by the
     `validate_on` keyword) which were provided on the last successful run.
 
     Args:
-        - state (State): a `Success` state from the last successful Task run which contains the cache
-        - inputs (dict): a `dict` of inputs which were available on the last
-            successful run of the cached Task
-        - parameters (dict): a `dict` of parameters which were available on the
-            last successful run of the cached Task
         - validate_on (list): a `list` of strings specifying the input names
             to validate against
 
     Returns:
-        - boolean specifying whether or not the cache should be used
+        - Callable: the actual validation function specifying whether or not the cache should be used
 
     Example:
     ```python
@@ -236,17 +236,40 @@ def rand_bool():
     state2 = f.run(task_states={result: state1.result[ans]})
     ```
     """
-    inputs = inputs or {}
-    if duration_only(state, inputs, parameters) is False:
-        return False
-    elif validate_on is None:
-        return True  # if you dont want to validate on anything, then the cache is valid
-    else:
-        cached = state.cached_inputs or {}
-        partial_provided = {
-            key: value for key, value in inputs.items() if key in validate_on
-        }
-        partial_needed = {
-            key: value for key, value in cached.items() if key in validate_on
-        }
-        return partial_provided == partial_needed
+
+    def _partial_inputs_only(
+        state: "prefect.engine.state.Cached",
+        inputs: Dict[str, Any],
+        parameters: Dict[str, Any],
+    ) -> bool:
+        """
+        The actual cache validation function which will be used.
+
+        Args:
+            - state (State): a `Success` state from the last successful Task run which contains the cache
+            - inputs (dict): a `dict` of inputs which were available on the last
+                successful run of the cached Task
+            - parameters (dict): a `dict` of parameters which were available on the
+                last successful run of the cached Task
+
+        Returns:
+            - boolean specifying whether or not the cache should be used
+        """
+        inputs = inputs or {}
+        if duration_only(state, inputs, parameters) is False:
+            return False
+        elif validate_on is None:
+            return (
+                True
+            )  # if you dont want to validate on anything, then the cache is valid
+        else:
+            cached = state.cached_inputs or {}
+            partial_provided = {
+                key: value for key, value in inputs.items() if key in validate_on
+            }
+            partial_needed = {
+                key: value for key, value in cached.items() if key in validate_on
+            }
+            return partial_provided == partial_needed
+
+    return _partial_inputs_only
diff --git a/src/prefect/serialization/task.py b/src/prefect/serialization/task.py
--- a/src/prefect/serialization/task.py
+++ b/src/prefect/serialization/task.py
@@ -16,6 +16,7 @@
 from prefect.utilities.serialization import (
     UUID,
     FunctionReference,
+    StatefulFunctionReference,
     JSONCompatible,
     ObjectSchema,
     from_qualified_name,
@@ -87,7 +88,7 @@ class Meta:
     inputs = fields.Method("load_inputs", allow_none=True)
     outputs = fields.Method("load_outputs", allow_none=True)
     timeout = fields.Integer(allow_none=True)
-    trigger = FunctionReference(
+    trigger = StatefulFunctionReference(
         valid_functions=[
             prefect.triggers.all_finished,
             prefect.triggers.manual_only,
@@ -105,7 +106,7 @@ class Meta:
     )
     skip_on_upstream_skip = fields.Boolean(allow_none=True)
     cache_for = fields.TimeDelta(allow_none=True)
-    cache_validator = FunctionReference(
+    cache_validator = StatefulFunctionReference(
         valid_functions=[
             prefect.engine.cache_validators.never_use,
             prefect.engine.cache_validators.duration_only,
diff --git a/src/prefect/utilities/serialization.py b/src/prefect/utilities/serialization.py
--- a/src/prefect/utilities/serialization.py
+++ b/src/prefect/utilities/serialization.py
@@ -1,4 +1,5 @@
 import base64
+import inspect
 import json
 import sys
 import types
@@ -330,3 +331,66 @@ def _deserialize(self, value, attr, data, **kwargs):  # type: ignore
         elif value not in self.valid_functions and self.reject_invalid:
             raise ValidationError("Invalid function reference: {}".format(value))
         return self.valid_functions.get(value, value)
+
+
+class StatefulFunctionReference(fields.Field):
+    """
+    Similar to `FunctionField`, Field that stores a reference to a possibly stateful function as a string and reloads it when
+    deserialized.
+
+    Args:
+        - valid_functions (List[Callable]): a list of functions that will be serialized as string
+            references
+        - reject_invalid (bool): if True, functions not in `valid_functions` will be rejected. If False,
+            any value will be allowed, but only functions in `valid_functions` will be deserialized.
+        - **kwargs (Any): the keyword arguments accepted by `marshmallow.Field`
+
+    """
+
+    def __init__(
+        self, valid_functions: list, reject_invalid: bool = True, **kwargs: Any
+    ):
+        self.valid_functions = {to_qualified_name(f): f for f in valid_functions}
+        self.reject_invalid = reject_invalid
+        super().__init__(**kwargs)
+
+    def _serialize(self, value, attr, obj, **kwargs):  # type: ignore
+        if value is None and self.allow_none:
+            return None
+
+        try:
+            qual_name = to_qualified_name(value)
+        except:
+            if self.reject_invalid:
+                raise ValidationError("Invalid function reference: {}".format(value))
+            else:
+                return qual_name
+
+        valid_bases = [
+            name for name in self.valid_functions if qual_name.startswith(name)
+        ]
+        if not any(valid_bases) and self.reject_invalid:
+            raise ValidationError("Invalid function reference: {}".format(value))
+        else:
+            base_name = next(filter(None, valid_bases)) if valid_bases else qual_name
+
+        state_kwargs = inspect.getclosurevars(value)[0]
+        call_sig = ", ".join(
+            ["{k}={v}".format(k=k, v=v) for k, v in state_kwargs.items()]
+        )
+        return base_name + "({})".format(call_sig)
+
+    def _deserialize(self, value, attr, data, **kwargs):  # type: ignore
+        """
+        Attempts to return the original function that was serialized, with _no state_.
+
+        If not found, returns `None`.
+        """
+        if value is None and self.allow_none:
+            return None
+
+        base_name = value.split("(")[0]
+        if base_name not in self.valid_functions and self.reject_invalid:
+            raise ValidationError("Invalid function reference: {}".format(value))
+
+        return self.valid_functions.get(base_name)


Test patch:
diff --git a/tests/engine/test_cache_validators.py b/tests/engine/test_cache_validators.py
--- a/tests/engine/test_cache_validators.py
+++ b/tests/engine/test_cache_validators.py
@@ -13,14 +13,8 @@
 )
 from prefect.engine.state import Cached
 
-all_validators = [
-    all_inputs,
-    all_parameters,
-    never_use,
-    partial_inputs_only,
-    duration_only,
-    partial_parameters_only,
-]
+all_validators = [all_inputs, all_parameters, never_use, duration_only]
+stateful_validators = [partial_inputs_only, partial_parameters_only]
 
 
 def test_never_use_returns_false():
@@ -33,6 +27,12 @@ def test_expired_cache(validator):
     assert validator(state, None, None) is False
 
 
+@pytest.mark.parametrize("validator", stateful_validators)
+def test_expired_cache_stateful(validator):
+    state = Cached(cached_result_expiration=pendulum.now("utc") - timedelta(days=1))
+    assert validator()(state, None, None) is False
+
+
 class TestDurationOnly:
     def test_unexpired_cache(self):
         state = Cached(cached_result_expiration=pendulum.now("utc") + timedelta(days=1))
@@ -74,37 +74,37 @@ def test_additional_parameters_invalidate(self):
 class TestPartialInputsOnly:
     def test_inputs_validate_with_defaults(self):
         state = Cached(cached_inputs=dict(x=1, s="str"))
-        assert partial_inputs_only(state, dict(x=1, s="str"), None) is True
+        assert partial_inputs_only(None)(state, dict(x=1, s="str"), None) is True
         state = Cached(cached_inputs=dict(x=1, s="str"))
-        assert partial_inputs_only(state, dict(x=1, s="strs"), None) is True
+        assert partial_inputs_only(None)(state, dict(x=1, s="strs"), None) is True
 
     def test_validate_on_kwarg(self):
         state = Cached(cached_inputs=dict(x=1, s="str"))
         assert (
-            partial_inputs_only(state, dict(x=1, s="str"), None, validate_on=["x", "s"])
+            partial_inputs_only(validate_on=["x", "s"])(state, dict(x=1, s="str"), None)
             is True
         )
         state = Cached(cached_inputs=dict(x=1, s="str"))
         assert (
-            partial_inputs_only(
-                state, dict(x=1, s="strs"), None, validate_on=["x", "s"]
+            partial_inputs_only(validate_on=["x", "s"])(
+                state, dict(x=1, s="strs"), None
             )
             is False
         )
         assert (
-            partial_inputs_only(state, dict(x=1, s="strs"), None, validate_on=["x"])
+            partial_inputs_only(validate_on=["x"])(state, dict(x=1, s="strs"), None)
             is True
         )
         assert (
-            partial_inputs_only(state, dict(x=1, s="strs"), None, validate_on=["s"])
+            partial_inputs_only(validate_on=["s"])(state, dict(x=1, s="strs"), None)
             is False
         )
 
     def test_handles_none(self):
         state = Cached(cached_parameters=dict(x=5))
-        assert partial_inputs_only(state, dict(x=5), None, validate_on=["x"]) is False
+        assert partial_inputs_only(validate_on=["x"])(state, dict(x=5), None) is False
         state = Cached(cached_inputs=dict(x=5))
-        assert partial_inputs_only(state, None, None, validate_on=["x"]) is False
+        assert partial_inputs_only(validate_on=["x"])(state, None, None) is False
 
     def test_curried(self):
         state = Cached(cached_inputs=dict(x=1, s="str"))
@@ -116,41 +116,41 @@ def test_curried(self):
 class TestPartialParametersOnly:
     def test_parameters_validate_with_defaults(self):
         state = Cached(cached_parameters=dict(x=1, s="str"))
-        assert partial_parameters_only(state, None, dict(x=1, s="str")) is True
+        assert partial_parameters_only()(state, None, dict(x=1, s="str")) is True
         state = Cached(cached_parameters=dict(x=1, s="str"))
-        assert partial_parameters_only(state, None, dict(x=1, s="strs")) is True
+        assert partial_parameters_only()(state, None, dict(x=1, s="strs")) is True
 
     def test_validate_on_kwarg(self):
         state = Cached(cached_parameters=dict(x=1, s="str"))
         assert (
-            partial_parameters_only(
-                state, None, dict(x=1, s="str"), validate_on=["x", "s"]
+            partial_parameters_only(validate_on=["x", "s"])(
+                state, None, dict(x=1, s="str")
             )
             is True
         )
         state = Cached(cached_parameters=dict(x=1, s="str"))
         assert (
-            partial_parameters_only(
-                state, None, dict(x=1, s="strs"), validate_on=["x", "s"]
+            partial_parameters_only(validate_on=["x", "s"])(
+                state, None, dict(x=1, s="strs")
             )
             is False
         )
         assert (
-            partial_parameters_only(state, None, dict(x=1, s="strs"), validate_on=["x"])
+            partial_parameters_only(validate_on=["x"])(state, None, dict(x=1, s="strs"))
             is True
         )
         assert (
-            partial_parameters_only(state, None, dict(x=1, s="strs"), validate_on=["s"])
+            partial_parameters_only(validate_on=["s"])(state, None, dict(x=1, s="strs"))
             is False
         )
 
     def test_handles_none(self):
         state = Cached(cached_inputs=dict(x=5))
         assert (
-            partial_parameters_only(state, None, dict(x=5), validate_on=["x"]) is False
+            partial_parameters_only(validate_on=["x"])(state, None, dict(x=5)) is False
         )
         state = Cached(cached_parameters=dict(x=5))
-        assert partial_parameters_only(state, None, None, validate_on=["x"]) is False
+        assert partial_parameters_only(validate_on=["x"])(state, None, None) is False
 
     def test_curried(self):
         state = Cached(cached_parameters=dict(x=1, s="str"))
diff --git a/tests/serialization/test_tasks.py b/tests/serialization/test_tasks.py
--- a/tests/serialization/test_tasks.py
+++ b/tests/serialization/test_tasks.py
@@ -123,14 +123,30 @@ def test_trigger(trigger):
     assert t2.trigger is trigger
 
 
+@pytest.mark.parametrize(
+    "trigger", [prefect.triggers.some_failed, prefect.triggers.some_successful]
+)
+@pytest.mark.parametrize(
+    "bounds", [(1, 5), (0.1, 0.9), (1, None), (0.1, None), (None, 42), (None, 0.5)]
+)
+def test_stateful_trigger(trigger, bounds):
+    kwargs = dict(zip(("at_least", "at_most"), bounds))
+    t = Task(trigger=trigger(**kwargs))
+    serialized = TaskSchema().dump(t)
+    call_sig = ", ".join(["{k}={v}".format(k=k, v=v) for k, v in kwargs.items()])
+    assert serialized["trigger"].endswith("({})".format(call_sig))
+
+    t2 = TaskSchema().load(serialized)
+    assert t2.trigger is trigger  # no state
+
+
 def test_unknown_trigger():
     def hello():
         pass
 
     t = Task(trigger=hello)
     t2 = TaskSchema().load(TaskSchema().dump(t))
-    assert isinstance(t2.trigger, str)
-    assert t2.trigger.endswith("hello")
+    assert t2.trigger is prefect.triggers.all_successful  # falls back to default
 
 
 @pytest.mark.parametrize(
@@ -157,16 +173,36 @@ def test_cache_validator_never_use():
     assert t2.cache_validator is prefect.engine.cache_validators.never_use
 
 
+@pytest.mark.parametrize(
+    "validator",
+    [
+        prefect.engine.cache_validators.partial_inputs_only,
+        prefect.engine.cache_validators.partial_parameters_only,
+    ],
+)
+@pytest.mark.parametrize("validate_on", [["x"], ["longer"], ["x", "y"]])
+def test_stateful_trigger(validator, validate_on):
+    with pytest.warns(UserWarning):
+        t = Task(cache_validator=validator(validate_on))
+    serialized = TaskSchema().dump(t)
+    call_sig = "validate_on={}".format(validate_on)
+    assert serialized["cache_validator"].endswith("({})".format(call_sig))
+
+    with pytest.warns(UserWarning):
+        t2 = TaskSchema().load(serialized)
+    assert t2.cache_validator is validator  # no state
+
+
 def test_unknown_cache_validator():
     def hello():
         pass
 
     with pytest.warns(UserWarning):
         t = Task(cache_validator=hello)
-    with pytest.warns(UserWarning):
-        t2 = TaskSchema().load(TaskSchema().dump(t))
-    assert isinstance(t2.cache_validator, str)
-    assert t2.cache_validator.endswith("hello")
+    t2 = TaskSchema().load(TaskSchema().dump(t))
+
+    # falls back to default
+    assert t2.cache_validator is prefect.engine.cache_validators.never_use
 
 
 def test_inputs_outputs():
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
@@ -13,6 +13,7 @@
     Bytes,
     DateTimeTZ,
     FunctionReference,
+    StatefulFunctionReference,
     JSONCompatible,
     Nested,
     ObjectSchema,
@@ -213,6 +214,73 @@ def test_deserialize_none(self):
         assert self.Schema().load({"f_none": None})["f_none"] is None
 
 
+def outer(x, y, z):
+    def inner(q):
+        return x + y + z + q
+
+    return inner
+
+
+class TestStatefulFunctionReferenceField:
+    class Schema(marshmallow.Schema):
+        f = StatefulFunctionReference(valid_functions=[outer])
+        f_allow_invalid = StatefulFunctionReference(
+            valid_functions=[outer], reject_invalid=False
+        )
+        f_none = StatefulFunctionReference(valid_functions=[outer], allow_none=True)
+
+    def test_serialize_outer_no_state(self):
+        serialized = self.Schema().dump(dict(f=outer))
+        assert serialized["f"] == "tests.utilities.test_serialization.outer()"
+
+    def test_serialize_outer_with_state(self):
+        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)"
+        )
+
+    def test_serialize_invalid_fn(self):
+        with pytest.raises(marshmallow.ValidationError):
+            self.Schema().dump(dict(f=fn2))
+
+    def test_serialize_invalid_fn_without_validation(self):
+        serialized = self.Schema().dump(dict(f_allow_invalid=fn2))
+        assert (
+            serialized["f_allow_invalid"] == "tests.utilities.test_serialization.fn2()"
+        )
+
+    def test_deserialize_outer_no_state(self):
+        deserialized = self.Schema().load(self.Schema().dump(dict(f=outer)))
+        assert deserialized["f"] is outer
+
+    def test_deserialize_outer_with_state(self):
+        deserialized = self.Schema().load(
+            self.Schema().dump(dict(f=outer(x=1, y=2, z=99)))
+        )
+        assert deserialized["f"] is outer
+
+    def test_deserialize_invalid_fn(self):
+        with pytest.raises(marshmallow.ValidationError):
+            self.Schema().load({"f": "hello"})
+
+    def test_deserialize_invalid_fn_without_validation(self):
+        deserialized = self.Schema().load(
+            dict(f_allow_invalid="tests.utilities.test_serialization.fn2")
+        )
+        assert deserialized["f_allow_invalid"] is None
+
+    def test_serialize_none(self):
+        with pytest.raises(marshmallow.ValidationError):
+            self.Schema().dump({"f": None})
+        assert self.Schema().dump({"f_none": None})["f_none"] is None
+
+    def test_deserialize_none(self):
+        with pytest.raises(marshmallow.ValidationError):
+            self.Schema().load({"f": None})
+        assert self.Schema().load({"f_none": None})["f_none"] is None
+
+
 class TestObjectSchema:
     def test_schema_writes_version_to_serialized_object(self):
         class TestObject: