894 lines
35 KiB
Python
894 lines
35 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""
|
|
Unit tests for the NixlConnector metadata handshake.
|
|
|
|
NIXL itself is stubbed out, so these tests exercise the ZMQ control plane and
|
|
the payload normalisation rather than any actual RDMA transfer.
|
|
"""
|
|
|
|
import ctypes
|
|
import sys
|
|
import time
|
|
import types
|
|
|
|
import msgspec
|
|
import pytest
|
|
import torch
|
|
import zmq
|
|
|
|
from vllm_omni.data_entry_keys import HiddenStatesStruct, MetaStruct, OmniPayloadStruct
|
|
|
|
pytestmark = [pytest.mark.cpu, pytest.mark.parallel, pytest.mark.core_model]
|
|
|
|
PORT = 47431
|
|
|
|
|
|
def _wait_for_metadata(consumer, key, metadata=None, *, timeout=2.0):
|
|
"""Retry outside the connector's deliberately bounded metadata query."""
|
|
deadline = time.monotonic() + timeout
|
|
while time.monotonic() < deadline:
|
|
resolved = consumer._resolve_metadata(key, metadata)
|
|
if resolved is not None:
|
|
return resolved
|
|
time.sleep(0.01)
|
|
pytest.fail(f"Timed out waiting for metadata for {key!r}")
|
|
|
|
|
|
@pytest.mark.parametrize("misses", [0, 2])
|
|
def test_wait_for_metadata_retries_until_success(monkeypatch, misses):
|
|
expected = {"claim_id": "claim"}
|
|
forwarded = {"source_host": "127.0.0.1", "source_port": PORT}
|
|
calls = []
|
|
|
|
def resolve(key, metadata):
|
|
calls.append((key, metadata))
|
|
return None if len(calls) <= misses else expected
|
|
|
|
ticks = iter(range(10))
|
|
monkeypatch.setattr(time, "monotonic", lambda: next(ticks))
|
|
monkeypatch.setattr(time, "sleep", lambda _: None)
|
|
consumer = types.SimpleNamespace(_resolve_metadata=resolve)
|
|
assert _wait_for_metadata(consumer, "retry", forwarded, timeout=5) is expected
|
|
assert calls == [("retry", forwarded)] * (misses + 1)
|
|
|
|
|
|
def test_wait_for_metadata_stops_at_deadline(monkeypatch):
|
|
calls = []
|
|
ticks = iter([0.0, 0.0, 1.0, 2.0])
|
|
monkeypatch.setattr(time, "monotonic", lambda: next(ticks))
|
|
monkeypatch.setattr(time, "sleep", lambda _: None)
|
|
consumer = types.SimpleNamespace(_resolve_metadata=lambda *args: calls.append(args))
|
|
with pytest.raises(pytest.fail.Exception, match="Timed out waiting for metadata for 'missing'"):
|
|
_wait_for_metadata(consumer, "missing")
|
|
assert calls == [("missing", None)] * 2
|
|
|
|
|
|
@pytest.mark.parametrize("async_chunk", [False, True])
|
|
def test_three_stage_incoming_and_outgoing_endpoints(nixl_connector_cls, async_chunk):
|
|
from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec, OmniTransferConfig
|
|
from vllm_omni.distributed.omni_connectors.utils.initialization import resolve_connector_spec
|
|
from vllm_omni.engine.stage_init_utils import get_stage_connector_spec
|
|
|
|
config = OmniTransferConfig(
|
|
connectors={
|
|
("0", "1"): ConnectorSpec(name="NixlConnector", extra={"host": "127.0.0.1", "zmq_port": PORT}),
|
|
("1", "2"): ConnectorSpec(name="NixlConnector", extra={"host": "127.0.0.1", "zmq_port": PORT + 10}),
|
|
}
|
|
)
|
|
connectors = []
|
|
try:
|
|
for stage in range(3):
|
|
spec = get_stage_connector_spec(config, stage, async_chunk)
|
|
resolved = resolve_connector_spec(ConnectorSpec(**spec), stage_id=stage, role=spec["extra"]["role"])
|
|
connectors.append(nixl_connector_cls(resolved.extra))
|
|
first, middle, last = connectors
|
|
assert middle._sender_zmq_port == PORT
|
|
assert middle._serving_handshake
|
|
assert middle._zmq_port == PORT + 11
|
|
first.put("0", "1", "incoming", torch.ones(1))
|
|
incoming = _wait_for_metadata(middle, "incoming")
|
|
assert incoming["sender_zmq_port"] == PORT
|
|
assert incoming["generation"] == first._published["incoming"]["generation"]
|
|
middle.put("1", "2", "outgoing", torch.ones(1))
|
|
outgoing = _wait_for_metadata(last, "outgoing")
|
|
assert outgoing["sender_zmq_port"] == PORT + 11
|
|
assert outgoing["generation"] == middle._published["outgoing"]["generation"]
|
|
finally:
|
|
for connector in reversed(connectors):
|
|
for pending in connector._pending.values():
|
|
pending.claims.clear() # No DMA in this metadata-only probe.
|
|
connector.close()
|
|
|
|
|
|
@pytest.mark.parametrize("name", ["NixlConnector", "MooncakeTransferEngineConnector"])
|
|
def test_explicit_sender_port_is_not_derived(name):
|
|
from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
|
|
from vllm_omni.distributed.omni_connectors.utils.initialization import resolve_connector_spec
|
|
|
|
spec = ConnectorSpec(name=name, extra={"sender_zmq_port": 55000, "from_stage": 0})
|
|
resolved = resolve_connector_spec(spec, stage_id=1, role="receiver", replica_id=3, local_rank=2)
|
|
assert resolved.extra["sender_zmq_port"] == 55000
|
|
assert spec.extra == {"sender_zmq_port": 55000, "from_stage": 0}
|
|
spec.extra["zmq_port"] = "${UNUSED_NIXL_BASE_PORT}"
|
|
assert resolve_connector_spec(spec, stage_id=1, role="receiver").extra["sender_zmq_port"] == 55000
|
|
|
|
|
|
@pytest.mark.parametrize("need_recv_cache", [False, True])
|
|
@pytest.mark.parametrize("update_before_init", [False, True])
|
|
def test_manager_receiver_does_not_bind_shared_producer_port(producer, need_recv_cache, update_before_init):
|
|
"""Payload-only and KV receivers must dial, not bind, the incoming edge."""
|
|
from vllm_omni.distributed.omni_connectors.kv_transfer_manager import (
|
|
OmniKVCacheConfig,
|
|
OmniKVTransferManager,
|
|
)
|
|
|
|
manager = OmniKVTransferManager(
|
|
OmniKVCacheConfig(
|
|
connector_config={
|
|
"type": "NixlConnector",
|
|
"role": "receiver",
|
|
"host": "127.0.0.1",
|
|
"zmq_port": PORT,
|
|
"backends": ["UCX"],
|
|
},
|
|
from_stage="0",
|
|
to_stage="1",
|
|
stage_id=1,
|
|
need_recv_cache=need_recv_cache,
|
|
)
|
|
)
|
|
sender_info = {"host": producer.host, "zmq_port": producer._zmq_port}
|
|
if update_before_init:
|
|
manager.update_sender_info(sender_info)
|
|
receiver = manager.connector
|
|
try:
|
|
assert receiver is not None, "Receiver incorrectly tried to bind the producer's occupied port"
|
|
assert receiver._zmq_port is None
|
|
assert not receiver._serving_handshake
|
|
if not update_before_init:
|
|
assert receiver._sender_zmq_port is None
|
|
manager.update_sender_info(sender_info)
|
|
assert receiver._sender_zmq_port == producer._zmq_port
|
|
assert manager.config.connector_config["zmq_port"] == PORT # Do not mutate the shared specification.
|
|
success, _, _ = producer.put("0", "1", "manager-receive", torch.ones(1))
|
|
assert success
|
|
metadata = _wait_for_metadata(receiver, "manager-receive")
|
|
assert metadata["generation"] == producer._published["manager-receive"]["generation"]
|
|
receiver._notify_transfer_done("manager-receive", metadata)
|
|
assert not producer._pending
|
|
finally:
|
|
if receiver is not None:
|
|
receiver.close()
|
|
|
|
|
|
@pytest.mark.parametrize("need_recv_cache", [False, True])
|
|
def test_manager_receiver_preserves_explicit_standalone_sender(producer, need_recv_cache):
|
|
from vllm_omni.distributed.omni_connectors.kv_transfer_manager import (
|
|
OmniKVCacheConfig,
|
|
OmniKVTransferManager,
|
|
)
|
|
|
|
manager = OmniKVTransferManager(
|
|
OmniKVCacheConfig(
|
|
connector_config={
|
|
"type": "NixlConnector",
|
|
"role": "receiver",
|
|
"zmq_port": PORT,
|
|
"sender_host": producer.host,
|
|
"sender_zmq_port": producer._zmq_port,
|
|
},
|
|
need_recv_cache=need_recv_cache,
|
|
)
|
|
)
|
|
receiver = manager.connector
|
|
try:
|
|
assert receiver is not None
|
|
assert receiver._zmq_port is None and not receiver._serving_handshake
|
|
assert (receiver._sender_host, receiver._sender_zmq_port) == (producer.host, producer._zmq_port)
|
|
success, _, _ = producer.put("0", "1", "standalone-receive", torch.ones(1))
|
|
assert success
|
|
metadata = _wait_for_metadata(receiver, "standalone-receive")
|
|
assert metadata["generation"] == producer._published["standalone-receive"]["generation"]
|
|
receiver._notify_transfer_done("standalone-receive", metadata)
|
|
assert not producer._pending
|
|
finally:
|
|
if receiver is not None:
|
|
receiver.close()
|
|
|
|
|
|
def test_middle_stage_advertises_its_outgoing_replica_endpoint():
|
|
from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient
|
|
|
|
client = object.__new__(StageEngineCoreClient)
|
|
client.stage_id = 1
|
|
client.replica_id = 3
|
|
client.vllm_config = types.SimpleNamespace(
|
|
model_config=types.SimpleNamespace(
|
|
stage_connector_config={
|
|
"name": "NixlConnector",
|
|
"extra": {
|
|
"role": "receiver",
|
|
"host": "10.0.0.1",
|
|
"zmq_port": 47000,
|
|
"from_stage": 0,
|
|
"outgoing": {"host": "10.0.0.2", "zmq_port": 48000, "from_stage": 1},
|
|
},
|
|
}
|
|
)
|
|
)
|
|
assert client._build_payload_sender_info() == {"host": "10.0.0.2", "zmq_port": 51073}
|
|
|
|
|
|
@pytest.mark.parametrize("direct", [False, True])
|
|
@pytest.mark.usefixtures("reliable_claim_queries")
|
|
def test_claimed_source_survives_expiry_and_cleanup(producer, consumer, direct):
|
|
_, _, metadata = producer.put("0", "1", "claimed", torch.ones(1))
|
|
resolved = consumer._resolve_metadata("claimed", metadata if direct else None)
|
|
pending = producer._pending["claimed"]
|
|
pending.deadline = 0
|
|
producer._cleanup_expired_pending()
|
|
assert producer._pending.get("claimed") is pending
|
|
producer.cleanup("claimed")
|
|
assert producer._pending.get("claimed") is pending
|
|
assert producer._agent.registered
|
|
consumer._notify_transfer_done("claimed", resolved)
|
|
consumer._notify_transfer_done("claimed", resolved)
|
|
assert producer._pending == {}
|
|
assert producer._agent.registered == []
|
|
|
|
|
|
@pytest.mark.usefixtures("reliable_claim_queries")
|
|
def test_claims_are_generation_scoped_and_duplicate_ack_cannot_release_sibling(producer, consumer):
|
|
_, _, original = producer.put("0", "1", "owners", torch.ones(1))
|
|
first = consumer._resolve_metadata("owners", original)
|
|
second = consumer._resolve_metadata("owners", original)
|
|
assert first["claim_id"] != second["claim_id"]
|
|
assert not producer.put("0", "1", "owners", torch.zeros(1))[0]
|
|
consumer._notify_transfer_done("owners", first)
|
|
consumer._notify_transfer_done("owners", first)
|
|
assert len(producer._pending["owners"].claims) == 1
|
|
consumer._notify_transfer_done("owners", second)
|
|
assert not producer._pending
|
|
producer.put("0", "1", "owners", torch.zeros(1))
|
|
consumer._notify_transfer_done("owners", second)
|
|
assert "owners" in producer._pending
|
|
assert consumer._resolve_metadata("owners", original) is None
|
|
|
|
|
|
@pytest.mark.parametrize("get_metadata", [None, {"schema_version": 1, "tensor_specs": [], "descriptor_groups": []}])
|
|
@pytest.mark.usefixtures("reliable_claim_queries")
|
|
def test_abandoned_claim_survives_close_and_late_completion(producer, consumer, get_metadata):
|
|
_, _, metadata = producer.put("0", "1", "abandoned", torch.ones(1))
|
|
claimed = consumer._resolve_metadata("abandoned", metadata)
|
|
producer._pending["abandoned"].deadline = 0
|
|
producer.close()
|
|
assert producer._closing and not producer._closed
|
|
assert producer._agent.registered
|
|
assert producer._listener_thread.is_alive()
|
|
with pytest.raises(RuntimeError, match="closed"):
|
|
producer.put("0", "1", "new", torch.ones(1))
|
|
metrics = dict(producer._metrics)
|
|
with pytest.raises(RuntimeError, match="Cannot get data: NixlConnector is closed"):
|
|
producer.get("0", "1", "new", get_metadata)
|
|
assert producer._metrics == metrics
|
|
assert producer._pending["abandoned"].claims == {claimed["claim_id"]}
|
|
assert consumer._resolve_metadata("abandoned", metadata) is None
|
|
consumer._notify_transfer_done("abandoned", claimed)
|
|
assert not producer._pending
|
|
assert not producer._agent.registered
|
|
assert producer._listener_thread.is_alive()
|
|
producer.close()
|
|
assert not producer._agent.registered
|
|
assert producer._closed
|
|
|
|
|
|
@pytest.fixture
|
|
def copying_native_agent(consumer, monkeypatch):
|
|
"""Strict fake descriptors and actual CPU copies; not native NIXL evidence."""
|
|
calls = []
|
|
|
|
def descriptors(regions, memory_type):
|
|
assert regions and all(region[0] > 0 and region[1] > 0 for region in regions)
|
|
return regions
|
|
|
|
def prepare(operation, local, local_ids, remote, remote_ids):
|
|
assert operation == "READ"
|
|
assert local_ids == remote_ids == list(range(len(local)))
|
|
return list(zip(local, remote, strict=True))
|
|
|
|
def transfer(pairs):
|
|
calls.append(pairs)
|
|
for destination, source in pairs:
|
|
assert destination[1] == source[1] > 0
|
|
ctypes.memmove(destination[0], source[0], source[1])
|
|
|
|
agent = consumer._agent
|
|
monkeypatch.setattr(agent, "add_remote_agent", lambda metadata: "producer", raising=False)
|
|
monkeypatch.setattr(agent, "get_xfer_descs", descriptors, raising=False)
|
|
monkeypatch.setattr(agent, "prep_xfer_dlist", lambda agent, descs: descs, raising=False)
|
|
monkeypatch.setattr(agent, "make_prepped_xfer", prepare, raising=False)
|
|
monkeypatch.setattr(agent, "transfer", transfer, raising=False)
|
|
monkeypatch.setattr(agent, "check_xfer_state", lambda handle: "DONE", raising=False)
|
|
for method in ("release_xfer_handle", "release_dlist_handle", "remove_remote_agent"):
|
|
monkeypatch.setattr(agent, method, lambda handle: None, raising=False)
|
|
return calls
|
|
|
|
|
|
@pytest.mark.parametrize("case", ["empty", "all_empty", "mixed", "scalar", "structured_empty", "structured_mixed"])
|
|
@pytest.mark.usefixtures("reliable_claim_queries")
|
|
def test_zero_byte_leaves_roundtrip_without_native_descriptors(producer, consumer, copying_native_agent, case):
|
|
empty = torch.empty((2, 0, 3), dtype=torch.float64)
|
|
other = torch.empty((0,), dtype=torch.int64)
|
|
scalar = torch.tensor(7)
|
|
vector = torch.arange(6, dtype=torch.float32)
|
|
payload = {
|
|
"empty": empty,
|
|
"all_empty": [empty, other],
|
|
"mixed": [empty, scalar, other, vector],
|
|
"scalar": scalar,
|
|
"structured_empty": {"leaves": (empty, [other]), "label": "empty"},
|
|
"structured_mixed": {"leaves": (empty, [scalar, other, vector]), "label": "mixed"},
|
|
}[case]
|
|
ok, size, metadata = producer.put("0", "1", "empty-slots", payload)
|
|
assert ok
|
|
indices = [index for group in metadata["descriptor_groups"] for index in group["tensor_indices"]]
|
|
assert sorted(indices) == [i for i, spec in enumerate(metadata["tensor_specs"]) if spec["size"]]
|
|
actual, received_size = consumer.get("0", "1", "empty-slots", metadata)
|
|
if isinstance(payload, dict):
|
|
assert actual["label"] == payload["label"]
|
|
torch.testing.assert_close(actual["leaves"], payload["leaves"], rtol=0, atol=0)
|
|
else:
|
|
torch.testing.assert_close(actual, payload, rtol=0, atol=0)
|
|
assert size == received_size
|
|
if case in ("empty", "all_empty"):
|
|
assert not copying_native_agent
|
|
assert not producer._pending and not producer._agent.registered
|
|
assert not consumer._agent.registered
|
|
|
|
|
|
@pytest.mark.parametrize("direct", [False, True])
|
|
@pytest.mark.parametrize("outcome", ["done", "error", "timeout", "unknown"])
|
|
@pytest.mark.usefixtures("reliable_claim_queries")
|
|
def test_read_ownership_through_terminal_and_deferred_paths(
|
|
producer, consumer, copying_native_agent, monkeypatch, direct, outcome
|
|
):
|
|
import threading
|
|
|
|
_, _, metadata = producer.put("0", "1", "active-read", torch.ones(2))
|
|
entered, proceed = threading.Event(), threading.Event()
|
|
state = ["PROC"]
|
|
monkeypatch.setattr(consumer._agent, "check_xfer_state", lambda handle: state[0])
|
|
|
|
def wait(handle, key):
|
|
entered.set()
|
|
assert proceed.wait(5)
|
|
if outcome == "done":
|
|
state[0] = "DONE"
|
|
elif outcome == "error":
|
|
state[0] = "ERR"
|
|
raise RuntimeError("terminal error")
|
|
else:
|
|
if outcome == "unknown":
|
|
state[0] = "UNKNOWN"
|
|
raise TimeoutError("still owned")
|
|
|
|
monkeypatch.setattr(consumer, "_wait_for_transfer", wait)
|
|
worker = threading.Thread(target=consumer.get, args=("0", "1", "active-read", metadata if direct else None))
|
|
worker.start()
|
|
try:
|
|
assert entered.wait(5)
|
|
pending = producer._pending["active-read"]
|
|
pending.deadline = 0
|
|
producer._cleanup_expired_pending()
|
|
producer.cleanup("active-read")
|
|
assert producer._pending["active-read"] is pending
|
|
assert producer._agent.registered
|
|
finally:
|
|
proceed.set()
|
|
worker.join(5)
|
|
assert not worker.is_alive()
|
|
if outcome in ("timeout", "unknown"):
|
|
assert producer._agent.registered
|
|
assert consumer._deferred_transfers
|
|
consumer._reap_deferred_transfers()
|
|
assert producer._agent.registered
|
|
state[0] = "DONE"
|
|
deadline = time.monotonic() + 5
|
|
while producer._pending and time.monotonic() < deadline:
|
|
time.sleep(0.01)
|
|
assert not producer._pending
|
|
assert not producer._agent.registered
|
|
|
|
|
|
class _FakeNixlAgent:
|
|
def __init__(self, name, config=None):
|
|
self.name = name
|
|
self.config = config
|
|
self.registered = []
|
|
|
|
def get_reg_descs(self, regions, memory_type):
|
|
return ("reg", tuple(regions), memory_type)
|
|
|
|
def register_memory(self, descs, backends=None):
|
|
self.registered.append(descs)
|
|
|
|
def deregister_memory(self, descs):
|
|
if descs in self.registered:
|
|
self.registered.remove(descs)
|
|
|
|
def get_agent_metadata(self):
|
|
return f"agent-metadata-{self.name}".encode()
|
|
|
|
|
|
@pytest.fixture
|
|
def nixl_connector_cls(monkeypatch):
|
|
"""Import NixlConnector with vLLM's optional NIXL dependency stubbed out."""
|
|
stub = types.ModuleType("vllm.distributed.nixl_utils")
|
|
stub.NixlWrapper = _FakeNixlAgent # type: ignore[attr-defined]
|
|
config_module = types.ModuleType("test_nixl_config_api")
|
|
config_sync_type = types.SimpleNamespace(NIXL_THREAD_SYNC_STRICT="strict")
|
|
config_module.nixl_thread_sync_t = config_sync_type # type: ignore[attr-defined]
|
|
|
|
def agent_config(**kwargs):
|
|
return kwargs
|
|
|
|
agent_config.__module__ = config_module.__name__
|
|
config_module.nixl_agent_config = agent_config # type: ignore[attr-defined]
|
|
stub.nixl_agent_config = agent_config # type: ignore[attr-defined]
|
|
monkeypatch.setitem(sys.modules, "vllm.distributed.nixl_utils", stub)
|
|
monkeypatch.setitem(sys.modules, config_module.__name__, config_module)
|
|
sync_type = types.SimpleNamespace(NIXL_THREAD_SYNC_STRICT="wrong-public-enum")
|
|
nixl_module = types.ModuleType("nixl")
|
|
nixl_module.nixl_thread_sync_t = sync_type # type: ignore[attr-defined]
|
|
monkeypatch.setitem(sys.modules, "nixl", nixl_module)
|
|
|
|
from vllm_omni.distributed.omni_connectors.connectors.nixl_connector import NixlConnector
|
|
|
|
return NixlConnector
|
|
|
|
|
|
@pytest.fixture
|
|
def producer(nixl_connector_cls):
|
|
connector = nixl_connector_cls({"role": "sender", "host": "127.0.0.1", "zmq_port": PORT})
|
|
yield connector
|
|
# These control-plane tests do not submit DMA. Explicitly complete claims
|
|
# made by metadata-only probes before tearing down the shared port.
|
|
for key, pending in list(connector._pending.items()):
|
|
for claim in list(pending.claims):
|
|
from vllm_omni.distributed.omni_connectors.connectors.nixl_connector import _XFER_DONE_MSG
|
|
|
|
connector._handle_handshake_message(
|
|
_XFER_DONE_MSG
|
|
+ msgspec.msgpack.encode({"key": key, "generation": pending.generation, "claim_id": claim})
|
|
)
|
|
connector.close()
|
|
|
|
|
|
@pytest.fixture
|
|
def consumer(nixl_connector_cls):
|
|
connector = nixl_connector_cls({"role": "receiver", "sender_host": "127.0.0.1", "sender_zmq_port": PORT})
|
|
yield connector
|
|
connector.close()
|
|
|
|
|
|
@pytest.fixture
|
|
def reliable_claim_queries(producer, consumer, monkeypatch):
|
|
"""Ownership probes need exact claims, not claims retained after lost replies.
|
|
|
|
Exercise the real resolver, wire encoding and producer handler synchronously;
|
|
discovery tests separately cover bounded queries over real ZMQ sockets.
|
|
"""
|
|
import uuid
|
|
|
|
from vllm_omni.distributed.omni_connectors.connectors.nixl_connector import _GET_META_MSG, _META_NOT_FOUND
|
|
|
|
def query(key, host, port, *, generation=None):
|
|
assert (host, port) == (producer.host, producer._zmq_port)
|
|
request = {"key": key, "generation": generation, "claim_id": uuid.uuid4().hex}
|
|
reply = producer._handle_handshake_message(_GET_META_MSG + msgspec.msgpack.encode(request))
|
|
return None if reply == _META_NOT_FOUND else msgspec.msgpack.decode(reply)
|
|
|
|
monkeypatch.setattr(consumer, "_query_metadata_at", query)
|
|
|
|
|
|
def test_put_publishes_its_handshake_endpoint(producer):
|
|
ok, size, metadata = producer.put("0", "1", "req-0", torch.arange(8, dtype=torch.float32))
|
|
|
|
assert ok is True
|
|
assert size == 32
|
|
assert metadata["sender_host"] == "127.0.0.1"
|
|
assert metadata["sender_zmq_port"] == PORT
|
|
|
|
|
|
def test_agent_uses_strict_thread_synchronization(producer):
|
|
assert producer._agent.config["sync_mode"] == "strict"
|
|
|
|
|
|
def test_handshake_serves_metadata_when_caller_has_none(producer, consumer):
|
|
_, _, published = producer.put("0", "1", "req-1", torch.arange(4, dtype=torch.float32))
|
|
|
|
resolved = _wait_for_metadata(consumer, "req-1")
|
|
|
|
# msgpack has no tuple type, so region descriptors arrive as lists; get()
|
|
# re-tuples them before handing them to NIXL.
|
|
assert resolved.pop("claim_id")
|
|
assert resolved == msgspec.msgpack.decode(msgspec.msgpack.encode(published))
|
|
assert [tuple(region) for group in resolved["descriptor_groups"] for region in group["regions"]] == [
|
|
region for group in published["descriptor_groups"] for region in group["regions"]
|
|
]
|
|
|
|
|
|
def test_legacy_metadata_without_generation_skips_the_handshake(consumer):
|
|
"""Externally owned metadata without a generation needs no ownership claim."""
|
|
direct = {"schema_version": 1, "kind": "tensors", "tensor_specs": []}
|
|
|
|
assert consumer._resolve_metadata("req-2", direct) is direct
|
|
|
|
|
|
def test_source_endpoint_metadata_overrides_configured_sender(consumer, monkeypatch):
|
|
queried = []
|
|
expected = {"schema_version": 1, "kind": "tensors", "tensor_specs": []}
|
|
|
|
def query(get_key, host, port):
|
|
queried.append((get_key, host, port))
|
|
return expected
|
|
|
|
monkeypatch.setattr(consumer, "_query_metadata_at", query)
|
|
|
|
resolved = consumer._resolve_metadata(
|
|
"req-tp4-rank3",
|
|
{"source_host": "10.0.0.3", "source_port": PORT + 3 * 16},
|
|
)
|
|
|
|
assert resolved is expected
|
|
assert queried == [("req-tp4-rank3", "10.0.0.3", PORT + 3 * 16)]
|
|
|
|
|
|
def test_unknown_key_is_queried_once(producer, consumer, monkeypatch):
|
|
socket = consumer._get_req_socket(f"tcp://127.0.0.1:{PORT}")
|
|
send_count = 0
|
|
original_send = socket.send
|
|
|
|
def count_send(message):
|
|
nonlocal send_count
|
|
send_count += 1
|
|
return original_send(message)
|
|
|
|
monkeypatch.setattr(socket, "send", count_send)
|
|
|
|
assert consumer._resolve_metadata("never-published", None) is None
|
|
assert send_count == 1
|
|
assert socket.getsockopt(zmq.RCVTIMEO) == 10
|
|
|
|
|
|
@pytest.mark.usefixtures("reliable_claim_queries")
|
|
def test_transfer_done_releases_the_producer_buffer(producer, consumer):
|
|
_, _, metadata = producer.put("0", "1", "req-3", torch.arange(4, dtype=torch.float32))
|
|
assert producer._pending and producer._agent.registered
|
|
|
|
metadata = consumer._resolve_metadata("req-3", metadata)
|
|
consumer._notify_transfer_done("req-3", metadata)
|
|
|
|
assert producer._pending == {}
|
|
assert producer._published == {}
|
|
assert producer._agent.registered == []
|
|
|
|
|
|
def test_direct_metadata_has_ephemeral_ownership_endpoint(nixl_connector_cls):
|
|
connector = nixl_connector_cls({})
|
|
try:
|
|
assert connector._zmq_ctx is not None
|
|
assert connector._listener_thread is not None
|
|
_, _, metadata = connector.put("0", "1", "req-4", torch.zeros(2))
|
|
assert metadata["sender_host"]
|
|
assert metadata["sender_zmq_port"] > 0
|
|
finally:
|
|
connector.close()
|
|
|
|
|
|
def test_idle_producer_lease_expires_without_another_put(nixl_connector_cls):
|
|
connector = nixl_connector_cls({"role": "sender", "lease_seconds": 0.01})
|
|
try:
|
|
connector.put("0", "1", "req-expire", torch.zeros(2))
|
|
|
|
deadline = time.monotonic() + 1.0
|
|
while connector._pending and time.monotonic() < deadline:
|
|
time.sleep(0.01)
|
|
|
|
assert connector._pending == {}
|
|
assert connector._published == {}
|
|
assert connector._agent.registered == []
|
|
assert connector._registered_descs == []
|
|
finally:
|
|
connector.close()
|
|
|
|
|
|
def test_close_stops_lease_reaper(nixl_connector_cls):
|
|
connector = nixl_connector_cls({"role": "sender"})
|
|
lease_thread = connector._lease_thread
|
|
|
|
connector.close()
|
|
|
|
assert lease_thread is not None
|
|
assert not lease_thread.is_alive()
|
|
assert connector._lease_thread is None
|
|
|
|
|
|
def test_deferred_transfer_is_retained_while_active(nixl_connector_cls):
|
|
from vllm_omni.distributed.omni_connectors.connectors.nixl_connector import _DeferredTransfer
|
|
|
|
connector = nixl_connector_cls({"role": "receiver"})
|
|
released = []
|
|
connector._agent.check_xfer_state = lambda handle: "PROC"
|
|
connector._agent.release_xfer_handle = lambda handle: released.append(("handle", handle))
|
|
connector._agent.release_dlist_handle = lambda handle: released.append(("dlist", handle))
|
|
connector._agent.remove_remote_agent = lambda agent: released.append(("agent", agent))
|
|
transfer = _DeferredTransfer(
|
|
tensors=[torch.zeros(1)],
|
|
registrations=["registration"],
|
|
dlists=["local", "remote"],
|
|
handles=["transfer"],
|
|
remote_agent="producer",
|
|
)
|
|
connector._defer_transfer(transfer)
|
|
try:
|
|
connector._reap_deferred_transfers()
|
|
|
|
assert released == []
|
|
assert transfer in connector._deferred_transfers
|
|
assert transfer.tensors
|
|
finally:
|
|
connector._agent.check_xfer_state = lambda handle: "DONE"
|
|
connector.close()
|
|
|
|
|
|
def test_deferred_transfer_releases_exactly_once_after_done(nixl_connector_cls):
|
|
from vllm_omni.distributed.omni_connectors.connectors.nixl_connector import _DeferredTransfer
|
|
|
|
connector = nixl_connector_cls({"role": "receiver"})
|
|
released = []
|
|
connector._agent.check_xfer_state = lambda handle: "DONE"
|
|
connector._agent.release_xfer_handle = lambda handle: released.append(("handle", handle))
|
|
connector._agent.release_dlist_handle = lambda handle: released.append(("dlist", handle))
|
|
connector._agent.remove_remote_agent = lambda agent: released.append(("agent", agent))
|
|
connector._agent.deregister_memory = lambda descs: released.append(("registration", descs))
|
|
transfer = _DeferredTransfer(
|
|
tensors=[torch.zeros(1)],
|
|
registrations=["registration"],
|
|
dlists=["local", "remote"],
|
|
handles=["transfer"],
|
|
remote_agent="producer",
|
|
)
|
|
connector._defer_transfer(transfer)
|
|
try:
|
|
connector._reap_deferred_transfers()
|
|
connector._reap_deferred_transfers()
|
|
|
|
assert released == [
|
|
("handle", "transfer"),
|
|
("dlist", "local"),
|
|
("dlist", "remote"),
|
|
("agent", "producer"),
|
|
("registration", "registration"),
|
|
]
|
|
assert connector._deferred_transfers == []
|
|
assert transfer.tensors == []
|
|
finally:
|
|
connector.close()
|
|
|
|
|
|
def test_active_sibling_transfer_requires_deferred_ownership(nixl_connector_cls):
|
|
connector = nixl_connector_cls({"role": "receiver"})
|
|
connector._agent.check_xfer_state = lambda handle: {"failed": "ERR", "active": "PROC"}[handle]
|
|
try:
|
|
assert connector._transfers_may_be_active(["failed", "active"]) is True
|
|
assert connector._transfers_may_be_active(["failed"]) is False
|
|
finally:
|
|
connector.close()
|
|
|
|
|
|
def test_get_defers_complete_ownership_when_sibling_transfer_is_active(nixl_connector_cls):
|
|
connector = nixl_connector_cls({"role": "receiver"})
|
|
connector._agent.add_remote_agent = lambda metadata: "producer"
|
|
connector._agent.get_xfer_descs = lambda regions, memory_type: (regions, memory_type)
|
|
connector._agent.prep_xfer_dlist = lambda agent, descs: (agent, tuple(descs[0]))
|
|
handles = iter(["failed", "active"])
|
|
connector._agent.make_prepped_xfer = lambda *_args: next(handles)
|
|
connector._agent.transfer = lambda handle: None
|
|
connector._agent.check_xfer_state = lambda handle: {"failed": "ERR", "active": "PROC"}[handle]
|
|
released = []
|
|
connector._agent.release_xfer_handle = lambda handle: released.append(("handle", handle))
|
|
connector._agent.release_dlist_handle = lambda handle: released.append(("dlist", handle))
|
|
connector._agent.remove_remote_agent = lambda agent: released.append(("agent", agent))
|
|
connector._agent.deregister_memory = lambda descs: released.append(("registration", descs))
|
|
metadata = {
|
|
"schema_version": 1,
|
|
"kind": "tensors",
|
|
"agent_metadata": b"producer-metadata",
|
|
"tensor_specs": [
|
|
{"shape": [1], "dtype": "torch.float32", "device": "cpu", "size": 4},
|
|
{"shape": [1], "dtype": "torch.float32", "device": "cpu", "size": 4},
|
|
],
|
|
"descriptor_groups": [
|
|
{"memory_type": "DRAM", "tensor_indices": [0], "regions": [(1, 4, 0, "")]},
|
|
{"memory_type": "DRAM", "tensor_indices": [1], "regions": [(2, 4, 0, "")]},
|
|
],
|
|
"size": 8,
|
|
}
|
|
|
|
try:
|
|
assert connector.get("0", "1", "req-active-sibling", metadata) is None
|
|
assert released == []
|
|
assert len(connector._deferred_transfers) == 1
|
|
deferred = connector._deferred_transfers[0]
|
|
assert deferred.handles == ["failed", "active"]
|
|
assert len(deferred.registrations) == 2
|
|
assert len(deferred.dlists) == 4
|
|
assert deferred.remote_agent == "producer"
|
|
assert len(deferred.tensors) == 2
|
|
finally:
|
|
connector._agent.check_xfer_state = lambda handle: "DONE"
|
|
connector.close()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"payload,expected_kind",
|
|
[
|
|
(torch.zeros(4), "tensors"),
|
|
([torch.zeros(2), torch.ones(3)], "tensors"),
|
|
({"hidden": torch.zeros(4), "meta": {"token_role_ids": [1, 2]}}, "structured"),
|
|
({"prompt": "a cat", "steps": 8}, "object"),
|
|
],
|
|
)
|
|
def test_payload_kinds_round_trip_through_the_handshake(producer, consumer, payload, expected_kind):
|
|
_, _, published = producer.put("0", "1", "req-kind", payload)
|
|
assert published["kind"] == expected_kind
|
|
|
|
resolved = _wait_for_metadata(consumer, "req-kind")
|
|
|
|
assert resolved["kind"] == expected_kind
|
|
assert resolved["tensor_specs"] == published["tensor_specs"]
|
|
|
|
|
|
def test_structured_payload_groups_descriptors_by_memory_type(producer, monkeypatch):
|
|
monkeypatch.setattr(
|
|
producer,
|
|
"_resolve_memory_type",
|
|
lambda tensor: "DRAM" if tensor.dtype == torch.uint8 else "VRAM",
|
|
)
|
|
|
|
_, _, metadata = producer.put(
|
|
"0",
|
|
"1",
|
|
"req-mixed",
|
|
{"meta": "value", "hidden": torch.ones(4, dtype=torch.float32)},
|
|
)
|
|
|
|
assert metadata["schema_version"] == 1
|
|
assert metadata["descriptor_groups"] == [
|
|
{
|
|
"memory_type": "DRAM",
|
|
"tensor_indices": [0],
|
|
"regions": metadata["descriptor_groups"][0]["regions"],
|
|
},
|
|
{
|
|
"memory_type": "VRAM",
|
|
"tensor_indices": [1],
|
|
"regions": metadata["descriptor_groups"][1]["regions"],
|
|
},
|
|
]
|
|
assert len(producer._agent.registered) == 2
|
|
|
|
|
|
def test_omni_payload_struct_uses_structured_tensor_path(producer):
|
|
payload = OmniPayloadStruct(
|
|
hidden_states=HiddenStatesStruct(output=torch.ones(4, dtype=torch.float32)),
|
|
meta=MetaStruct(left_context_size=2),
|
|
)
|
|
|
|
_, _, metadata = producer.put("0", "1", "req-struct", payload)
|
|
skeleton, tensors = producer._extract_tensor_leaves(payload)
|
|
|
|
assert metadata["kind"] == "structured"
|
|
assert len(tensors) == 1
|
|
assert torch.equal(tensors[0], payload.hidden_states.output)
|
|
assert skeleton["meta"]["left_context_size"] == 2
|
|
|
|
|
|
def test_partial_group_registration_failure_rolls_back(producer, monkeypatch):
|
|
monkeypatch.setattr(
|
|
producer,
|
|
"_resolve_memory_type",
|
|
lambda tensor: "DRAM" if tensor.dtype == torch.uint8 else "VRAM",
|
|
)
|
|
original_register = producer._agent.register_memory
|
|
|
|
def fail_second_group(descs, backends=None):
|
|
if descs[2] == "VRAM":
|
|
raise RuntimeError("registration failed")
|
|
original_register(descs, backends=backends)
|
|
|
|
monkeypatch.setattr(producer._agent, "register_memory", fail_second_group)
|
|
|
|
success, _, metadata = producer.put(
|
|
"0",
|
|
"1",
|
|
"req-partial",
|
|
{"meta": "value", "hidden": torch.ones(4, dtype=torch.float32)},
|
|
)
|
|
|
|
assert success is False
|
|
assert metadata is None
|
|
assert producer._agent.registered == []
|
|
assert producer._registered_descs == []
|
|
assert "req-partial" not in producer._pending
|
|
|
|
|
|
def test_reusing_put_key_releases_previous_registration(producer):
|
|
producer.put("0", "1", "req-reuse", torch.zeros(2))
|
|
previous_registration = producer._agent.registered[0]
|
|
|
|
producer.put("0", "1", "req-reuse", torch.ones(2))
|
|
|
|
assert previous_registration not in producer._agent.registered
|
|
assert len(producer._agent.registered) == 1
|
|
assert len(producer._registered_descs) == 1
|
|
|
|
|
|
def test_expired_snapshot_cannot_claim_replacement(producer):
|
|
from vllm_omni.distributed.omni_connectors.connectors.nixl_connector import _PendingPayload
|
|
|
|
old = _PendingPayload([torch.zeros(1)], ["old"], 0.0)
|
|
replacement = _PendingPayload([torch.ones(1)], ["new"], time.monotonic() + 60)
|
|
producer._pending["req-race"] = replacement
|
|
|
|
assert producer._take_pending("req-race", expected=old) is None
|
|
assert producer._pending["req-race"] is replacement
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"indices",
|
|
[
|
|
[0, 0],
|
|
[0],
|
|
[0, 2],
|
|
[-1, 0],
|
|
],
|
|
)
|
|
def test_descriptor_groups_require_exact_tensor_index_partition(nixl_connector_cls, indices):
|
|
metadata = {
|
|
"schema_version": 1,
|
|
"descriptor_groups": [
|
|
{
|
|
"memory_type": "DRAM",
|
|
"tensor_indices": indices,
|
|
"regions": [(1, 4, 0, "")] * len(indices),
|
|
}
|
|
],
|
|
}
|
|
|
|
with pytest.raises(RuntimeError, match="exact partition"):
|
|
nixl_connector_cls._validated_descriptor_groups(
|
|
metadata, [{"shape": [1], "dtype": "torch.float32", "size": 4}] * 2
|
|
)
|
|
|
|
|
|
def test_receive_device_ignores_the_producer_index(nixl_connector_cls):
|
|
"""A producer on cuda:3 must not dictate the consumer's card."""
|
|
connector = nixl_connector_cls({})
|
|
try:
|
|
assert connector._resolve_receive_device("cpu") == torch.device("cpu")
|
|
assert connector._resolve_receive_device(None) == torch.device("cpu")
|
|
finally:
|
|
connector.close()
|
|
|
|
|
|
def test_receive_device_config_wins(nixl_connector_cls):
|
|
connector = nixl_connector_cls({"receive_device": "cpu"})
|
|
try:
|
|
assert connector._resolve_receive_device("cuda:3") == torch.device("cpu")
|
|
finally:
|
|
connector.close()
|