Skip to content

← Parallel Python with Process Pools step 15 of 24

Hard Research

Pickle a Large Buffer Without Copying It

You want to send a 100 MB bytearray to a worker. pickle.dumps copies all 100 MB into the pickle stream, and pickle.loads copies them again into a new object. Pickle protocol 5 can skip both copies. Wrap the data in pickle.PickleBuffer and pass a buffer_callback. Pickle then hands the buffer to your callback instead of writing its bytes into the stream. A buffer sent this way travels out of band, beside the stream instead of inside it. To load, give pickle.loads the same buffers in the same order.

import pickle

source = bytearray(b"abc")
buffers: list[pickle.PickleBuffer] = []
stream = pickle.dumps(
    pickle.PickleBuffer(source),
    protocol=5,
    buffer_callback=buffers.append,
)
view = memoryview(pickle.loads(stream, buffers=buffers))
source[0] = ord("z")
assert view.tobytes() == b"zbc"  # the view sees the change, so nothing was copied
assert view.obj is source

buffers.append returns None, and a callback that returns a false value tells pickle to keep that buffer out of band.

Between processes, something else must carry the buffers, such as shared memory. multiprocessing.Queue does not, so a bytearray on a queue still travels inside the stream. The grader cannot start processes, so you pack and unpack in one process. That is why the loaded payload can be the original bytearray.

Only unpickle data you trust, because pickle.loads can run code.

What to write

pack(tag, payload, protocol) returns (stream, buffers).

  • For protocol 5 or higher, pickle (tag, pickle.PickleBuffer(payload)) with a callback that collects the buffers in a list. Return the stream and the list.
  • For a lower protocol, pickle (tag, bytes(payload)) with no callback, and return the stream and an empty list. Older protocols cannot pickle a PickleBuffer.
  • Pass the requested protocol to pickle.dumps.

unpack(stream, buffers) loads the stream with those buffers and returns (tag, memoryview(raw_payload)), without copying the payload again.

solve(*, tag, pattern, repeat, protocol) builds source = bytearray(pattern.encode() * repeat), packs it, unpacks it with the returned buffers, and returns six values in this order:

  1. The restored tag.
  2. Whether len(stream) < 1024.
  3. How many buffers pack returned.
  4. Whether view.obj is source. True means the view uses the original memory.
  5. The payload length in bytes, len(view).
  6. The first 16 characters of hashlib.sha256(view).hexdigest().

protocol is between 0 and pickle.HIGHEST_PROTOCOL, which is 5, and repeat is zero or more. A character in pattern can take more than one byte in UTF-8.

At protocol 5 you get one buffer and view.obj is source, even for an empty payload. At lower protocols you get no buffers and a separate, read-only bytes object. A short payload gives a stream under 1 KiB at any protocol.