diff --git a/distributed/protocol/pickle.py b/distributed/protocol/pickle.py index 8b4b7328e5b..04b27174638 100644 --- a/distributed/protocol/pickle.py +++ b/distributed/protocol/pickle.py @@ -67,14 +67,19 @@ def dumps(x, *, buffer_callback=None, protocol=HIGHEST_PROTOCOL): buffers.clear() pickler.dump(x) result = f.getvalue() - if b"__main__" in result or ( + + if not _always_use_pickle_for(x) and ( CLOUDPICKLE_GE_20 and getattr(inspect.getmodule(x), "__name__", None) in cloudpickle.list_registry_pickle_by_value() + or ( + len(result) < 1000 + # Do this very last since it's expensive + and b"__main__" in result + ) ): - if len(result) < 1000 or not _always_use_pickle_for(x): - buffers.clear() - result = cloudpickle.dumps(x, **dump_kwargs) + buffers.clear() + result = cloudpickle.dumps(x, **dump_kwargs) except Exception: try: buffers.clear()