iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Bugfix][Distributed] Own SHM out-of-band buffers before releasing read slots (#61000)

Browse Source 
Signed-off-by: Hubert Marek <163992334+hubert-marek@users.noreply.github.com>
Signed-off-by: Ao Shen <aoshen@inferact.ai>
Co-authored-by: Codex <noreply@openai.com>
Co-authored-by: Ao Shen <aoshen@inferact.ai>
This commit is contained in:
Hubert
2026-10-12 01:44:12 -04:00
committed by GitHub
co-authored by Codex Ao Shen
parent fe4df3df0e
commit 74f4404f58
2 changed files with 87 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/distributed/test_shm_broadcast.py
+80 -25
View File
@@ -57,6 +57,7 @@ def distributed_run(fn, world_size, timeout=60, mp_context=None):
ctx = mp_context or mp
number_of_processes = world_size
processes = []
master_port = str(get_open_port())
for i in range(number_of_processes):
env = {}
env["RANK"] = str(i)
@@ -64,7 +65,7 @@ def distributed_run(fn, world_size, timeout=60, mp_context=None):
env["WORLD_SIZE"] = str(number_of_processes)
env["LOCAL_WORLD_SIZE"] = str(number_of_processes)
env["MASTER_ADDR"] = "localhost"
env["MASTER_PORT"] = "12345"
env["MASTER_PORT"] = master_port
p = ctx.Process(target=fn, args=(env,))
processes.append(p)
p.start()
@@ -88,11 +89,14 @@ def distributed_run(fn, world_size, timeout=60, mp_context=None):
break
time.sleep(0.1) # Check every 100ms
# Check for timeout if no failures detected yet
# Reap every child; failures can occur after the last polling iteration.
failed_processes.clear()
for i, p in enumerate(processes):
if p.is_alive():
p.kill()
p.join()
p.join()
if p.exitcode != 0:
failed_processes.append((i, p.exitcode))
# Report failures
if failed_processes:
@@ -102,6 +106,11 @@ def distributed_run(fn, world_size, timeout=60, mp_context=None):
raise AssertionError(error_msg)
def test_distributed_run_reports_timeout():
with pytest.raises(AssertionError, match="Distributed test failed"):
distributed_run(lambda _: time.sleep(60), world_size=1, timeout=0)
def worker_fn_wrapper(fn):
# `mp.Process` cannot accept environment variables directly
# so we need to pass the environment variables as arguments
@@ -381,14 +390,10 @@ def worker_fn_tensor_broadcast():
# Both ranks construct the identical reference payload.
torch.manual_seed(42)
payload = {
# 2MiB: rides the shm ring as an out-of-band buffer (the receiving
# side must copy out of the reusable ring chunk).
"mid": torch.randn(1024, 512),
# 16MiB > max_chunk_bytes: overflows to the zmq socket (the
# receiving side aliases the zmq.Frame zero-copy).
"big": torch.randn(4096, 2048, dtype=torch.bfloat16),
"nested": ["plain", 123, {"inner": torch.arange(5)}],
}
big = torch.randn(4096, 2048, dtype=torch.bfloat16)
if rank == writer_rank:
with mock.patch(
@@ -396,28 +401,23 @@ def worker_fn_tensor_broadcast():
wraps=_reduce_tensor,
) as wrapped_reduce:
message_queue.enqueue(payload)
# Overflow is per message, not per tensor: send this separately.
message_queue.enqueue(big)
assert wrapped_reduce.call_count == 3
# Cycle the ring (max_chunks=4) several times over so that aliased
# ring chunks would be overwritten.
for i in range(16):
message_queue.enqueue({"junk": torch.full((1024, 512), float(i))})
else:
received = message_queue.dequeue(timeout=30)
for key in ("mid", "big"):
assert torch.equal(received[key], payload[key]), key
assert received[key].dtype == payload[key].dtype, key
with mock.patch.object(MessageQueue, "recv", wraps=MessageQueue.recv) as recv:
received = message_queue.dequeue(timeout=30)
recv.assert_not_called()
received_big = message_queue.dequeue(timeout=30)
recv.assert_called_once()
assert torch.equal(received["mid"], payload["mid"])
assert received["mid"].dtype == payload["mid"].dtype
assert torch.equal(received_big, big)
assert received_big.dtype == big.dtype
assert torch.equal(received["nested"][2]["inner"], torch.arange(5))
snapshot = received["mid"].clone()
for i in range(16):
junk = message_queue.dequeue(timeout=30)
assert torch.equal(junk["junk"], torch.full((1024, 512), float(i)))
# Tensors received via the shm ring must not alias chunk memory
# that the writer has reused for subsequent messages.
assert torch.equal(received["mid"], snapshot)
# Rebuilt tensors must be writable, like regular tensors.
received["mid"] += 1.0
received["big"][0, 0] = 1.0
received_big[0, 0] = 1.0
dist.barrier()
print(f"tensor broadcast passed the test! Rank {rank}")
@@ -427,6 +427,61 @@ def test_tensor_broadcast():
distributed_run(worker_fn_tensor_broadcast, 2)
@pytest.mark.parametrize("kind", ["tensor", "numpy", "numpy_readonly"])
@pytest.mark.parametrize(
"nbytes, overflow, num_oob",
[
pytest.param(1024 * 1024 - 2, False, 0, id="inline"),
pytest.param(1024 * 1024, False, 1, id="shm_oob"),
pytest.param(2 * 1024 * 1024, True, 1, id="overflow"),
],
)
def test_broadcast_retains_values_after_ring_reuse(kind, nbytes, overflow, num_oob):
"""Received data belongs to its consumer, not the next queue message."""
writer = MessageQueue(
n_reader=1,
n_local_reader=1,
max_chunk_bytes=2 * 1024 * 1024,
max_chunks=1,
connect_ip="127.0.0.1",
)
reader = MessageQueue.create_from_handle(writer.export_handle(), rank=0)
writer.wait_until_ready()
reader.wait_until_ready()
array = np.full(nbytes // 2, 17, dtype=np.uint16)
array.flags.writeable = kind != "numpy_readonly"
payload = torch.from_numpy(array) if kind == "tensor" else array
try:
with (
mock.patch.object(MessageQueue, "recv", wraps=MessageQueue.recv) as recv,
mock.patch.object(pickle, "loads", wraps=pickle.loads) as loads,
):
writer.enqueue(payload)
received = reader.dequeue(timeout=5)
received_oob = len(loads.call_args.kwargs["buffers"])
# The spy must not keep the ring's memoryviews alive.
loads.reset_mock()
assert recv.call_count == int(overflow)
assert received_oob == num_oob
assert type(received) is type(payload)
assert received.dtype == payload.dtype
if isinstance(received, np.ndarray):
assert received.flags.writeable == array.flags.writeable
np.testing.assert_array_equal(received, array)
for value in range(4):
# Fit the replacement in SHM even when the original overflowed.
writer.enqueue(np.full(512 * 1024, value, dtype=np.uint16))
replacement = reader.dequeue(timeout=5)
np.testing.assert_array_equal(replacement, value)
np.testing.assert_array_equal(received, array)
if array.flags.writeable:
received[0] = 42
assert array[0] == 17
finally:
reader.shutdown()
writer.shutdown()
def _dumps_oob(obj) -> tuple[bytes, list]:
"""Pickle `obj` the same way `MessageQueue.enqueue` does: tensor
dispatch table + out-of-band buffers >= 1MiB."""
vllm/distributed/device_communicators/shm_broadcast.py
+7 -8
View File
@@ -408,8 +408,8 @@ def _rebuild_tensor(buf: Any, shape: tuple[int, ...], dtype_str: str) -> torch.T
Counterpart of `_reduce_tensor`. Note that pickle passes the original
buffer-providing object from `loads(buffers=...)` straight to this
function (no `PickleBuffer` wrapper on the receiving side), so `buf` is
a `zmq.Frame`, a `memoryview` of a shared-memory ring chunk, or `bytes`
if the buffer was serialized in-band.
a `zmq.Frame`, an owned `bytearray` from the shared-memory ring, or
`bytes`/`bytearray` if the buffer was serialized in-band.
"""
dtype = getattr(torch, dtype_str)
assert isinstance(dtype, torch.dtype)
@@ -423,11 +423,8 @@ def _rebuild_tensor(buf: Any, shape: tuple[int, ...], dtype_str: str) -> torch.T
except ValueError:
# Empty or read-only frame buffer; fall through to the copy path.
pass
# Shared-memory ring buffer chunks are reused by the writer once all
# readers have marked them read, so we must copy out of them. bytearray
# (vs bytes) keeps the resulting tensor writable, matching normal tensor
# semantics.
raw = bytearray(buf)
# dequeue already owns SHM buffers. Copy other buffers to keep them writable.
raw = buf if isinstance(buf, bytearray) else bytearray(buf)
if not raw:
assert 0 in shape
return torch.empty(shape, dtype=dtype)
@@ -981,7 +978,9 @@ class MessageQueue:
buf_offset = offset + 4
buf_len = from_bytes_big(buf[offset:buf_offset])
offset = buf_offset + buf_len
all_buffers.append(buf[buf_offset:offset])
buffer = buf[buf_offset:offset]
# Out-of-band arrays may outlive this reusable ring slot.
all_buffers.append(buffer if i == 0 else bytearray(buffer))
obj = pickle.loads(all_buffers[0], buffers=all_buffers[1:])
if overflow:
obj = MessageQueue.recv(self.local_socket, timeout)