vllm-omni/tests/distributed/omni_connectors/test_nixl_connector.py

# 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"])
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"])
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()