Allow setting of task cache_policy globally

[Jezreel-Zamora-Paidy]

Hello,
We recently run into an error when migrating to Prefect version 3 where one of our tasks failed to compute the cache key of a tasks because the input parameter contains Kafka DeserializingConsumer object and it fails to pickle it ( See the error below). Most of our tasks does not really need task caching and a lot of them contains an input parameter that is similar DeserializingConsumer object which will fail with the similar error with the default cache_policy.

Having an ability to set the cache_policy globally should be an option so that we don't have to set the cache_policy for every tasks that we have. Also if we are not setting PREFECT_RESULTS_PERSIST_BY_DEFAULT to True, does it still make sense for the tasks to calculate the cache key?

Error message

Error encountered when computing cache key - result will not be persisted.
Traceback (most recent call last):
  File "/opt/venv/lib/python3.11/site-packages/prefect/task_engine.py", line 155, in compute_transaction_key
    key = self.task.cache_policy.compute_key(
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/venv/lib/python3.11/site-packages/prefect/cache_policies.py", line 168, in compute_key
    policy_key = policy.compute_key(
                 ^^^^^^^^^^^^^^^^^^^
  File "/opt/venv/lib/python3.11/site-packages/prefect/cache_policies.py", line 168, in compute_key
    policy_key = policy.compute_key(
                 ^^^^^^^^^^^^^^^^^^^
  File "/opt/venv/lib/python3.11/site-packages/prefect/cache_policies.py", line 296, in compute_key
    return hash_objects(hashed_inputs, raise_on_failure=True)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/venv/lib/python3.11/site-packages/prefect/utilities/hashing.py", line 66, in hash_objects
    return stable_hash(cloudpickle.dumps((args, kwargs)), hash_algo=hash_algo)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/venv/lib/python3.11/site-packages/cloudpickle/cloudpickle.py", line 1529, in dumps
    cp.dump(obj)
  File "/opt/venv/lib/python3.11/site-packages/cloudpickle/cloudpickle.py", line 1295, in dump
    return super().dump(obj)
           ^^^^^^^^^^^^^^^^^
TypeError: cannot pickle 'DeserializingConsumer' object

[desertaxle]

Hey @Jezreel-Zamora-Paidy!

I think we might be able to handle this better by failing more gracefully when cache key computation fails. We could fall back to the NONE cache policy and log a warning in cases like the one you mentioned. We could also add an option to suppress that warning if it's too noisy. I think it'd be good to try that before introducing a default cache policy setting since it would require no action on the user's part. Would the fallback approach work for your use case?

[Jezreel-Zamora-Paidy]

Hey @desertaxle,

Thank you for looking on to this. I think falling back the cache_policy to NONE should work for most of our use case.

Is there any performance issues with using the default cache_policy with a tasks with Iterable as an input parameter?
Before encountering the issue mentioned above, our flow started failing with OOM error when running a tasks with an Iterable as an input parameter. After we increase the memory of the flow run, we encounter the error mentioned in this discussion. Setting the cache policy to None for that task also fixed the OOM error.

I think the computation of cache_key for tasks with Iterable parameter is consuming too much Memory.

[desertaxle]

Interesting, do you have an example where you see excessive memory consumption?

[Jezreel-Zamora-Paidy]

We encounter the OOM error when running a tasks called upload_stream_to_s3 which uploads the passed data to an s3 bucket.
As I cannot give the full source code, here is the pseudocode.

def upload_stream_to_s3(data: Iterable,
                     batch_size: int = 50_000,
                     s3_uri_prefix: Optional[str] = None,
                     use_naive_jst_timestamps: bool = True):
                     
                     # loop on data
                         # Add to buffer until batch_size >=  buffer
                         # upload buffer to s3                        
                     # return S3Uri

Let me know if you need more information.