From 701dbbdef919d5a403c487648021c1a35c40fbfc Mon Sep 17 00:00:00 2001 From: fjetter Date: Fri, 12 Jan 2024 14:00:46 +0100 Subject: [PATCH] Revert "Do not always check if `__main__ in result` when pickling (#8443)" This reverts commit 5c481ddcbf77ce814797d663068d6c0651844caf. --- distributed/protocol/pickle.py | 13 ++++--------- 1 file changed, 4 insertions(+), 9 deletions(-) diff --git a/distributed/protocol/pickle.py b/distributed/protocol/pickle.py index 04b27174638..8b4b7328e5b 100644 --- a/distributed/protocol/pickle.py +++ b/distributed/protocol/pickle.py @@ -67,19 +67,14 @@ def dumps(x, *, buffer_callback=None, protocol=HIGHEST_PROTOCOL): buffers.clear() pickler.dump(x) result = f.getvalue() - - if not _always_use_pickle_for(x) and ( + if b"__main__" in result or ( 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 - ) ): - buffers.clear() - result = cloudpickle.dumps(x, **dump_kwargs) + if len(result) < 1000 or not _always_use_pickle_for(x): + buffers.clear() + result = cloudpickle.dumps(x, **dump_kwargs) except Exception: try: buffers.clear()