Allow pickle protocol to be overridden

distributed/distributed-schema.yaml

@@ -759,6 +759,18 @@ properties:
             type: boolean
             description: Enter Python Debugger on scheduling error
 
+  pickle:
+    type: object
+    description: |
+      Configuration for pickle serialization
+    properties:
+      protocol:
+        type:
+        - integer
+        - "null"
+        description:
+          The protocol version to use with pickle
+
   rmm:
     type: object
     description: |

distributed/distributed.yaml

@@ -170,6 +170,8 @@ distributed:
     log-length: 10000  # default length of logs to keep in memory
     log-format: '%(name)s - %(levelname)s - %(message)s'
     pdb-on-err: False       # enter debug mode on scheduling error
+pickle:
+  protocol: null  # specify the pickle protocol to use
 rmm:
   pool-size: null
 ucx:

distributed/protocol/pickle.py

@@ -2,6 +2,7 @@
 import sys
 
 import cloudpickle
+import dask
 
 if sys.version_info < (3, 8):
     try:
@@ -33,16 +34,17 @@ def _always_use_pickle_for(x):
         return False
 
 
-def dumps(x, *, buffer_callback=None):
+def dumps(x, protocol=None, *, buffer_callback=None):
     """ Manage between cloudpickle and pickle
 
     1.  Try pickle
     2.  If it is short then check if it contains __main__
     3.  If it is long, then first check type, then check __main__
     """
+    protocol = protocol or dask.config.get("pickle.protocol") or HIGHEST_PROTOCOL
     buffers = []
-    dump_kwargs = {"protocol": HIGHEST_PROTOCOL}
-    if HIGHEST_PROTOCOL >= 5 and buffer_callback is not None:
+    dump_kwargs = {"protocol": protocol}
+    if protocol >= 5 and buffer_callback is not None:
         dump_kwargs["buffer_callback"] = buffers.append
     try:
         buffers.clear()

distributed/protocol/tests/test_pickle.py

@@ -25,6 +25,13 @@ def test_pickle_data():
         assert deserialize(*serialize(d, serializers=("pickle",))) == d
 
 
+def test_pickle_protocol():
+    data = {"int": 1, "float": 2, "unicode": "abc", "bytes": b"def", "set": set()}
+    for p in range(HIGHEST_PROTOCOL):
+        assert loads(dumps(data, p)) == data
+        assert deserialize(*serialize(data, serializers=("pickle",))) == data
+
+
 def test_pickle_out_of_band():
     class MemoryviewHolder:
         def __init__(self, mv):