Repository: PrefectHQ/prefect

Issue:
Make context pickleable
## Current behavior
In the following flow:

```python
from prefect import context, Flow, task
from prefect.engine.executors import DaskExecutor

@task
def extract():
    """Get a list of data"""
    return [1, 2, 3]


@task
def transform(data):
    """Multiply the input by 10"""
    return [i * 10 for i in data]


@task
def load(data):
    """Print the data to indicate it was received"""
    logger = context.get("logger")
    logger.info("Here's your data: {}".format(data))

with Flow("ETL") as flow:
    e = extract()
    t = transform(e)
    l = load(t)

executor = DaskExecutor(address="localhost:8786"
flow.run(executor=executor)
```

We will see the error `Unexpected error occured in FlowRunner: TypeError("can't pickle Context objects")` because `context` is used as an attribute of `prefect`. This requires users to `import prefect` in their tasks in order to use context.


## Proposed behavior
We could adjust the behavior so that context would be unpickled in the task and return `prefect.context` therefore avoiding this extra step of having to import prefect.



Discussion:
This is a good idea and in principle is trivial to implement.  However, all implementations I've tried so far preserve context as it was at task creation, which means that all of the things we put into context during a task run (e.g., `logger`) aren't available.   

The reason for this is that the serialization occurs at task submission time, and the deserialization occurs prior to the task runner actually running anything, so there's never really a chance to update context with the appropriate information. 

I think instead of making context pickleable, for now I'm going to simply raise an informative error message.

Resolving patch:
diff --git a/src/prefect/utilities/context.py b/src/prefect/utilities/context.py
--- a/src/prefect/utilities/context.py
+++ b/src/prefect/utilities/context.py
@@ -80,6 +80,21 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
             self.update(config.context)
         self["config"] = merge_dicts(config, self.get("config", {}))  # order matters
 
+    def __getstate__(self) -> None:
+        """
+        Because we dynamically update context during runs, we don't ever want to pickle
+        or "freeze" the contents of context.  Consequently it should always be accessed
+        as an attribute of the prefect module.
+        """
+        raise TypeError(
+            "\n".join(
+                [
+                    "Pickling context objects is explicitly not supported.",
+                    "You should always access context as an attribute of the `prefect` module, as in `prefect.context`",
+                ]
+            )
+        )
+
     def __repr__(self) -> str:
         return "<Context>"
 


Test patch:
diff --git a/tests/test_context.py b/tests/test_context.py
--- a/tests/test_context.py
+++ b/tests/test_context.py
@@ -1,3 +1,4 @@
+import cloudpickle
 import queue
 import threading
 import time
@@ -181,3 +182,10 @@ def get_context_in_thread(q, id, delay=0):
         results.add(result_queue.get(block=False))
 
     assert results == set(range(len(threads)))
+
+
+def test_context_raises_informative_error_if_pickled():
+    c = Context()
+
+    with pytest.raises(TypeError, match="prefect.context"):
+        cloudpickle.dumps(c)