| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import msgspec |
| 4 | |
| 5 | from vllm.distributed.ec_transfer.ec_connector.cpu.protocol import ( |
| 6 | EC_CONNECTOR_VERSION, |
| 7 | XferAck, |
| 8 | XferReq, |
| 9 | XferStatus, |
| 10 | compute_ec_compatibility_hash, |
| 11 | ) |
| 12 | |
| 13 | |
| 14 | def test_xferreq_roundtrip(): |
| 15 | req = XferReq(mm_hash="h1", compatibility_hash="c1", session_id="s1") |
| 16 | data = msgspec.msgpack.encode(req) |
| 17 | out = msgspec.msgpack.decode(data, type=XferReq) |
| 18 | assert out.mm_hash == "h1" |
| 19 | assert out.compatibility_hash == "c1" |
| 20 | assert out.connector_version == EC_CONNECTOR_VERSION |
| 21 | |
| 22 | |
| 23 | def test_xferack_roundtrip_ok(): |
| 24 | ack = XferAck( |
| 25 | mm_hash="h1", |
| 26 | status=XferStatus.OK, |
| 27 | src_block_indices=[1, 2], |
| 28 | agent_metadata=b"meta", |
| 29 | mem_descriptor=b"desc", |
| 30 | ) |
| 31 | out = msgspec.msgpack.decode(msgspec.msgpack.encode(ack), type=XferAck) |
| 32 | assert out.status == XferStatus.OK |
| 33 | assert out.src_block_indices == [1, 2] |
| 34 | |
| 35 | |
| 36 | def test_compat_hash_deterministic_and_sensitive(): |
| 37 | a = compute_ec_compatibility_hash("0.1", "m", "float16", 64) |
| 38 | b = compute_ec_compatibility_hash("0.1", "m", "float16", 64) |
| 39 | c = compute_ec_compatibility_hash("0.1", "m", "float16", 128) |
| 40 | assert a == b |
| 41 | assert a != c |