revert encode_into changes

Signed-off-by: Staszek Pasko <[email protected]>

Author: p88h

tests/v1/test_serial_utils.py

@@ -70,12 +70,14 @@ def test_encode_decode():
 
     assert_equal(decoded, obj)
 
-    # Test whether MsgpackEncoder properly reuses the buffers.
+    # Test encode_into case
 
-    encoded2 = encoder.encode(obj)
+    preallocated = bytearray()
+
+    encoded2 = encoder.encode_into(obj, preallocated)
 
     assert len(encoded2) == 6
-    assert encoded2[0] is encoded[0]
+    assert encoded2[0] is preallocated
 
     decoded2: MyType = decoder.decode(encoded2)
 

vllm/v1/engine/core.py

@@ -507,15 +507,15 @@ def process_output_socket(self, output_path: str, engine_index: int):
         """Output socket IO thread."""
 
         # Msgpack serialization encoding.
-        # The wrapper keeps an internal encoding buffer that avoids
-        # creating a new buffer for each encode call.
         encoder = MsgpackEncoder()
+        # Reuse send buffer.
+        buffer = bytearray()
 
         with zmq_socket_ctx(output_path, zmq.constants.PUSH) as socket:
             while True:
                 outputs = self.output_queue.get()
                 outputs.engine_index = engine_index
-                buffers = encoder.encode(outputs)
+                buffers = encoder.encode_into(outputs, buffer)
                 socket.send_multipart(buffers, copy=False)
 
 

vllm/v1/serial_utils.py

@@ -58,19 +58,27 @@ def __init__(self, size_threshold: Optional[int] = None):
         # This is used as a local stash of buffers that we can then access from
         # our custom `msgspec` hook, `enc_hook`. We don't have a way to
         # pass custom data to the hook otherwise.
-        self.msg_buffer = bytearray()
         self.aux_buffers: Optional[list[bytestr]] = None
         self.size_threshold = size_threshold
 
     def encode(self, obj: Any) -> Sequence[bytestr]:
         try:
+            self.aux_buffers = bufs = [b'']
+            bufs[0] = self.encoder.encode(obj)
             # This `bufs` list allows us to collect direct pointers to backing
             # buffers of tensors and np arrays, and return them along with the
             # top-level encoded buffer instead of copying their data into the
             # new buffer.
-            self.aux_buffers = [self.msg_buffer]
-            bufs = self.aux_buffers
-            self.encoder.encode_into(obj, self.msg_buffer)
+            return bufs
+        finally:
+            self.aux_buffers = None
+
+    # TODO: would be nice to make this automatic
+    def encode_into(self, obj: Any, buf: bytearray) -> Sequence[bytestr]:
+        try:
+            self.aux_buffers = [buf]
+            bufs = [buf]
+            self.encoder.encode_into(obj, buf)
             return bufs
         finally:
             self.aux_buffers = None