Revert pickle change (#8456)

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()

distributed/protocol/tests/test_pickle.py

@@ -20,7 +20,7 @@
     loads,
 )
 from distributed.protocol.serialize import dask_deserialize, dask_serialize
-from distributed.utils_test import save_sys_modules
+from distributed.utils_test import popen, save_sys_modules
 
 
 class MemoryviewHolder:
@@ -278,3 +278,22 @@ def test_nopickle_nested():
     finally:
         del dask_serialize._lookup[NoPickle]
         del dask_deserialize._lookup[NoPickle]
+
+
+@pytest.mark.slow()
+def test_pickle_functions_in_main(tmp_path):
+    script = """
+from dask.distributed import Client
+if __name__ == "__main__":
+    with Client(n_workers=1) as client:
+        def func(df):
+            return (df + 5)
+        client.submit(func, 5).result()
+        print("script successful", flush=True)
+"""
+    with open(tmp_path / "script.py", mode="w") as f:
+        f.write(script)
+    with popen([sys.executable, tmp_path / "script.py"], capture_output=True) as proc:
+        out, _ = proc.communicate(timeout=60)
+
+    assert "script successful" in out.decode("utf-8")