From 5c481ddcbf77ce814797d663068d6c0651844caf Mon Sep 17 00:00:00 2001 From: Florian Jetter Date: Wed, 10 Jan 2024 12:14:38 +0100 Subject: [PATCH] Do not always check if `__main__ in result` when pickling (#8443) --- distributed/protocol/pickle.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/distributed/protocol/pickle.py b/distributed/protocol/pickle.py index 8b4b7328e5..04b2717463 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()