py-sharedmemory 4.2.0__tar.gz → 4.2.1__tar.gz

{py_sharedmemory-4.2.0 → py_sharedmemory-4.2.1}/CHANGELOG.md

@@ -1,5 +1,15 @@
 # Changelog
 
+## 4.2.1
+
+- `put` copies large `bytes` and `bytearray` objects (64 KiB and up) straight into shared
+  memory, wherever they appear in a message. It used to build the complete pickle in
+  memory first: an extra copy of the payload that briefly needed up to 1.5 times its size
+  (48 MiB for a 32 MiB `bytearray`). Strided NumPy arrays, which NumPy pickles as `bytes`,
+  skip that copy too. The message format is unchanged.
+- Resizing a `bytearray` while its message is being pickled, from a `__reduce__` method,
+  now raises `BufferError` instead of risking a corrupted message.
+
 ## 4.2.0
 
 - Reuse received shared-memory blocks instead of freeing them. Writing into fresh pages

{py_sharedmemory-4.2.0/py_sharedmemory.egg-info → py_sharedmemory-4.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: py-sharedmemory
-Version: 4.2.0
+Version: 4.2.1
 Summary: Fast multiprocessing.Queue alternative that moves large Python objects and NumPy arrays through shared memory
 Author: Tobias Würth
 License-Expression: MIT
@@ -68,8 +68,9 @@ images, and simulation states arrive up to 2× faster, and transfers too large f
 
 - **Familiar API:** `put` / `get` with `block`, `timeout`, `put_nowait` / `get_nowait`,
   `queue.Full`, and `queue.Empty`.
-- **Any picklable object:** NumPy arrays and other buffers use pickle protocol 5
-  out-of-band buffers, so they are copied straight into shared memory.
+- **Any picklable object:** NumPy arrays, other protocol-5 buffers, and large `bytes` and
+  `bytearray` objects are copied straight into shared memory, without an intermediate
+  pickle copy.
 - **Bounded memory:** a capacity limit applies backpressure, and received blocks are
   reused for later messages, so shared memory stays within `capacity × message size`.
 - **One producer, many consumers:** consumer processes and threads share one receiver.
@@ -161,9 +162,10 @@ pages on every `put`.
 
 ## How it works
 
-1. `put()` pickles the object with protocol 5. Large buffers, such as NumPy array data,
-   stay out-of-band, and everything is copied into one shared-memory block: a block from
-   an earlier, received message when one fits, otherwise a new one.
+1. `put()` pickles the object with protocol 5. Large buffers, such as NumPy array data and
+   `bytes` or `bytearray` objects of 64 KiB or more, never enter the pickle; everything is copied once,
+   from the original objects into one shared-memory block: a block from an earlier,
+   received message when one fits, otherwise a new one.
 2. Only the block's name, buffer sizes, and a sequence number travel through a
    `multiprocessing.Queue`.
 3. `get()` attaches to the block, copies the data into local memory, unpickles it, and

{py_sharedmemory-4.2.0 → py_sharedmemory-4.2.1}/README.md

@@ -18,8 +18,9 @@ images, and simulation states arrive up to 2× faster, and transfers too large f
 
 - **Familiar API:** `put` / `get` with `block`, `timeout`, `put_nowait` / `get_nowait`,
   `queue.Full`, and `queue.Empty`.
-- **Any picklable object:** NumPy arrays and other buffers use pickle protocol 5
-  out-of-band buffers, so they are copied straight into shared memory.
+- **Any picklable object:** NumPy arrays, other protocol-5 buffers, and large `bytes` and
+  `bytearray` objects are copied straight into shared memory, without an intermediate
+  pickle copy.
 - **Bounded memory:** a capacity limit applies backpressure, and received blocks are
   reused for later messages, so shared memory stays within `capacity × message size`.
 - **One producer, many consumers:** consumer processes and threads share one receiver.
@@ -111,9 +112,10 @@ pages on every `put`.
 
 ## How it works
 
-1. `put()` pickles the object with protocol 5. Large buffers, such as NumPy array data,
-   stay out-of-band, and everything is copied into one shared-memory block: a block from
-   an earlier, received message when one fits, otherwise a new one.
+1. `put()` pickles the object with protocol 5. Large buffers, such as NumPy array data and
+   `bytes` or `bytearray` objects of 64 KiB or more, never enter the pickle; everything is copied once,
+   from the original objects into one shared-memory block: a block from an earlier,
+   received message when one fits, otherwise a new one.
 2. Only the block's name, buffer sizes, and a sequence number travel through a
    `multiprocessing.Queue`.
 3. `get()` attaches to the block, copies the data into local memory, unpickles it, and

{py_sharedmemory-4.2.0 → py_sharedmemory-4.2.1}/py_sharedmemory/__init__.py

@@ -25,5 +25,5 @@ from .factory import create_shared_memory_pair
 from .receiver import SharedMemoryReceiver
 from .sender import SharedMemorySender
 
-__version__ = "4.2.0"
+__version__ = "4.2.1"
 __all__ = ["SharedMemoryReceiver", "SharedMemorySender", "__version__", "create_shared_memory_pair"]

{py_sharedmemory-4.2.0 → py_sharedmemory-4.2.1}/py_sharedmemory/convert.py

@@ -37,6 +37,16 @@ def release_shared_memory(allocation: SharedMemory) -> None:
             allocation.unlink()
 
 
+class _PickleChunks(list[memoryview]):
+    """Pickle output kept as views of what pickle wrote rather than one joined copy."""
+
+    def write(self, data: Any) -> None:
+        # Pickle passes bytes and bytearrays of 64 KiB or more as the objects themselves.
+        # A view also stops a bytearray from being resized before it is copied, which would
+        # leave its length header in the stream wrong.
+        self.append(memoryview(data))
+
+
 def serialize(
     data: Any,
     sequence: int = 0,
@@ -44,8 +54,11 @@ def serialize(
 ) -> tuple[SharedMemory, SharedMemoryMessage]:
     """Copy a pickle and its out-of-band buffers into an owned block of at least that size.
 
-    `allocate` returns a new or reused block; this function frees it if copying fails.
+    Large bytes, bytearrays, and contiguous arrays are copied once, straight from the
+    original objects into shared memory. `allocate` returns a new or reused block; this
+    function frees it if copying fails.
     """
+    chunks = _PickleChunks()
     buffers: list[memoryview] = []
     allocation = None
 
@@ -53,23 +66,23 @@ def serialize(
         buffers.append(buffer.raw())
 
     try:
-        pickle_data = pickle.dumps(data, protocol=5, buffer_callback=append_buffer)
-        buffer_sizes = (len(pickle_data), *(buffer.nbytes for buffer in buffers))
+        pickle.Pickler(chunks, protocol=5, buffer_callback=append_buffer).dump(data)
+        pickle_size = sum(chunk.nbytes for chunk in chunks)
+        buffer_sizes = (pickle_size, *(buffer.nbytes for buffer in buffers))
         allocation = allocate(sum(buffer_sizes))
         shared = cast(memoryview, allocation.buf)
-        shared[: len(pickle_data)] = pickle_data
-        offset = len(pickle_data)
-        for buffer in buffers:
-            shared[offset : offset + buffer.nbytes] = buffer
-            offset += buffer.nbytes
+        offset = 0
+        for view in (*chunks, *buffers):
+            shared[offset : offset + view.nbytes] = view
+            offset += view.nbytes
         return allocation, SharedMemoryMessage(allocation.name, buffer_sizes, sequence)
     except BaseException:
         if allocation is not None:
             release_shared_memory(allocation)
         raise
     finally:
-        for buffer in buffers:
-            buffer.release()
+        for view in (*chunks, *buffers):
+            view.release()
 
 
 def deserialize(message: SharedMemoryMessage) -> Any:

{py_sharedmemory-4.2.0 → py_sharedmemory-4.2.1}/py_sharedmemory/sender.py

@@ -104,8 +104,9 @@ class SharedMemorySender:
         """Send a picklable object through shared memory.
 
         Args:
-            data: Any object that pickle protocol 5 can serialize. NumPy arrays and
-                other buffers are copied out-of-band without an extra pickle copy.
+            data: Any object that pickle protocol 5 can serialize. NumPy arrays, other
+                buffers, and large bytes-like objects are copied without an extra
+                pickle copy.
             block: Wait for capacity when `capacity` messages are outstanding.
             timeout: Maximum seconds to wait for capacity; `None` waits forever.
 

{py_sharedmemory-4.2.0 → py_sharedmemory-4.2.1/py_sharedmemory.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: py-sharedmemory
-Version: 4.2.0
+Version: 4.2.1
 Summary: Fast multiprocessing.Queue alternative that moves large Python objects and NumPy arrays through shared memory
 Author: Tobias Würth
 License-Expression: MIT
@@ -68,8 +68,9 @@ images, and simulation states arrive up to 2× faster, and transfers too large f
 
 - **Familiar API:** `put` / `get` with `block`, `timeout`, `put_nowait` / `get_nowait`,
   `queue.Full`, and `queue.Empty`.
-- **Any picklable object:** NumPy arrays and other buffers use pickle protocol 5
-  out-of-band buffers, so they are copied straight into shared memory.
+- **Any picklable object:** NumPy arrays, other protocol-5 buffers, and large `bytes` and
+  `bytearray` objects are copied straight into shared memory, without an intermediate
+  pickle copy.
 - **Bounded memory:** a capacity limit applies backpressure, and received blocks are
   reused for later messages, so shared memory stays within `capacity × message size`.
 - **One producer, many consumers:** consumer processes and threads share one receiver.
@@ -161,9 +162,10 @@ pages on every `put`.
 
 ## How it works
 
-1. `put()` pickles the object with protocol 5. Large buffers, such as NumPy array data,
-   stay out-of-band, and everything is copied into one shared-memory block: a block from
-   an earlier, received message when one fits, otherwise a new one.
+1. `put()` pickles the object with protocol 5. Large buffers, such as NumPy array data and
+   `bytes` or `bytearray` objects of 64 KiB or more, never enter the pickle; everything is copied once,
+   from the original objects into one shared-memory block: a block from an earlier,
+   received message when one fits, otherwise a new one.
 2. Only the block's name, buffer sizes, and a sequence number travel through a
    `multiprocessing.Queue`.
 3. `get()` attaches to the block, copies the data into local memory, unpickles it, and

{py_sharedmemory-4.2.0 → py_sharedmemory-4.2.1}/tests/test_convert.py

@@ -1,5 +1,6 @@
 import pickle
 import sys
+import tracemalloc
 from multiprocessing.shared_memory import SharedMemory
 from unittest.mock import Mock
 
@@ -155,26 +156,106 @@ def test_protocol_5_buffer_mutability_and_empty_buffers(payload):
         assert payload[0] != 0
 
 
-def test_pickle_failure_releases_exported_views():
-    payload = bytearray(b"data")
-    with pytest.raises(ValueError, match="Cannot serialize"):
-        serialize(FailsAfterBuffer(payload))
-    # Resizing would raise BufferError if the raw pickle buffer were retained.
+MiB = 2**20
+# Pickle hands bytes-like objects from 64 KiB up to its output file by reference.
+LARGE = 2**17
+
+
+@pytest.mark.parametrize(
+    "payload",
+    [
+        pytest.param(lambda: bytes(32 * MiB), id="bytes"),
+        pytest.param(lambda: bytearray(32 * MiB), id="bytearray"),
+        pytest.param(
+            lambda: {"step": 7, "frames": bytearray(32 * MiB), "tags": ["train"] * 100},
+            id="nested",
+        ),
+    ],
+)
+def test_large_bytes_reach_shared_memory_without_an_intermediate_copy(payload):
+    original = payload()
+    tracemalloc.start()
+    try:
+        handle, info = serialize(original)
+        peak = tracemalloc.get_traced_memory()[1]
+    finally:
+        tracemalloc.stop()
+    try:
+        # pickle.dumps would allocate the whole payload again, and more while growing.
+        assert peak < MiB
+        assert deserialize(info) == original
+    finally:
+        release_shared_memory(handle)
+
+
+class ResizesEarlierPayload:
+    def __init__(self, payload):
+        self.payload = payload
+
+    def __reduce__(self):
+        self.payload.extend(b"more")
+        return int, ()
+
+
+def test_bytearray_cannot_be_resized_before_it_is_copied():
+    # A resize after pickle wrote the length header would corrupt the stream.
+    payload = bytearray(LARGE)
+    with pytest.raises(BufferError):
+        serialize([payload, ResizesEarlierPayload(payload)])
     payload.extend(b"more")
+    assert len(payload) == LARGE + 4
 
 
-def test_allocation_failure_releases_exported_views(monkeypatch):
-    payload = bytearray(b"data")
+@pytest.mark.parametrize(
+    ("payload", "wrap"),
+    [
+        pytest.param(bytearray(b"data"), FailsAfterBuffer, id="out-of-band"),
+        pytest.param(bytearray(LARGE), lambda payload: [payload, Unpicklable()], id="in-band"),
+    ],
+)
+def test_pickle_failure_releases_exported_views(payload, wrap):
+    with pytest.raises(ValueError, match="Cannot serialize"):
+        serialize(wrap(payload))
+    # Resizing would raise BufferError if a view of the payload were retained.
+    payload.extend(b"more")
+
 
+@pytest.mark.parametrize(
+    ("payload", "wrap"),
+    [
+        pytest.param(bytearray(b"data"), Buffered, id="out-of-band"),
+        pytest.param(bytearray(LARGE), list, id="in-band"),
+    ],
+)
+def test_allocation_failure_releases_exported_views(monkeypatch, payload, wrap):
     def fail_allocation(**kwargs):
         raise MemoryError("No shared memory")
 
     monkeypatch.setattr(convert, "SharedMemory", fail_allocation)
     with pytest.raises(MemoryError):
-        serialize(Buffered(payload))
+        serialize(wrap(payload))
     payload.extend(b"more")
 
 
+def test_success_releases_exported_views():
+    payload = bytearray(LARGE)
+    handle, info = serialize({"payload": payload})
+    try:
+        payload.extend(b"more")
+        assert deserialize(info) == {"payload": bytearray(LARGE)}
+    finally:
+        release_shared_memory(handle)
+
+
+def test_payload_spanning_many_pickle_frames_round_trips():
+    original = {"numbers": list(range(200_000)), "blob": b"x" * LARGE, "text": "y" * LARGE}
+    handle, info = serialize(original)
+    try:
+        assert deserialize(info) == original
+    finally:
+        release_shared_memory(handle)
+
+
 def test_copy_failure_unlinks_new_allocation(monkeypatch):
     created = []