| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Unit tests for ECCPUConnector scheduler utilities.""" |
| 4 | |
| 5 | import msgspec |
| 6 | import pytest |
| 7 | |
| 8 | from vllm.distributed.ec_transfer.ec_connector.cpu.utils import ( |
| 9 | build_block_descs, |
| 10 | deserialize_mem_descriptor, |
| 11 | serialize_mem_descriptor, |
| 12 | ) |
| 13 | |
| 14 | # ── build_block_descs ──────────────────────────────────────────────────────── |
| 15 | |
| 16 | |
| 17 | @pytest.mark.parametrize("device_id,expected_dev", [(7, 7), (None, 0)]) |
| 18 | def test_build_block_descs(device_id, expected_dev): |
| 19 | kwargs = dict(base_ptr=1000, num_blocks=4, block_size_bytes=256) |
| 20 | if device_id is not None: |
| 21 | kwargs["device_id"] = device_id |
| 22 | descs = build_block_descs(**kwargs) |
| 23 | assert len(descs) == 4 |
| 24 | for i, (addr, size, dev) in enumerate(descs): |
| 25 | assert addr == 1000 + i * 256 |
| 26 | assert size == 256 |
| 27 | assert dev == expected_dev |
| 28 | |
| 29 | |
| 30 | def test_build_block_descs_zero_blocks_returns_empty(): |
| 31 | assert build_block_descs(base_ptr=100, num_blocks=0, block_size_bytes=64) == [] |
| 32 | |
| 33 | |
| 34 | # ── serialize / deserialize_mem_descriptor ─────────────────────────────────── |
| 35 | |
| 36 | |
| 37 | def test_mem_descriptor_roundtrip(): |
| 38 | descs = [(100, 64, 0), (164, 64, 0), (228, 64, 1)] |
| 39 | assert deserialize_mem_descriptor(serialize_mem_descriptor(descs)) == descs |
| 40 | |
| 41 | |
| 42 | @pytest.mark.parametrize( |
| 43 | "bad_value", |
| 44 | [ |
| 45 | [(1, 2)], # 2-tuple instead of 3 |
| 46 | [("a", "b", "c")], # strings, not ints |
| 47 | ], |
| 48 | ) |
| 49 | def test_mem_descriptor_rejects_malformed_payload(bad_value): |
| 50 | """Malformed descriptor lists must fail to decode.""" |
| 51 | encoder = msgspec.msgpack.Encoder() |
| 52 | bad_payload = encoder.encode(bad_value) |
| 53 | with pytest.raises((msgspec.DecodeError, msgspec.ValidationError)): |
| 54 | deserialize_mem_descriptor(bad_payload) |