102 lines
3.0 KiB
Python
102 lines
3.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from vllm.config import ParallelConfig
|
|
from vllm.v1.executor import vllm_net_devices
|
|
from vllm.v1.executor.vllm_net_devices import normalize_pci
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"addr, expected",
|
|
[
|
|
("0000:3f:00.0", (0, 0x3F, 0, 0)),
|
|
("0001:00:00.0", (1, 0, 0, 0)),
|
|
("00000001:00:00.0", (1, 0, 0, 0)),
|
|
("0009:00:00.0", (9, 0, 0, 0)),
|
|
("0105:00:00.0", (0x105, 0, 0, 0)),
|
|
("0000:0a:1f.7", (0, 0x0A, 0x1F, 7)),
|
|
],
|
|
)
|
|
def test_normalize_pci_full_domain(addr, expected):
|
|
assert normalize_pci(addr) == expected
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"addr, expected",
|
|
[
|
|
("01:00.0", (0, 1, 0, 0)),
|
|
("3f:00.0", (0, 0x3F, 0, 0)),
|
|
("40:00.0", (0, 0x40, 0, 0)),
|
|
("ff:1f.7", (0, 0xFF, 0x1F, 7)),
|
|
],
|
|
)
|
|
def test_normalize_pci_short_form(addr, expected):
|
|
assert normalize_pci(addr) == expected
|
|
|
|
|
|
def test_normalize_pci_case_insensitive():
|
|
assert normalize_pci("0A:1F.7") == normalize_pci("0a:1f.7")
|
|
|
|
|
|
def test_normalize_pci_strips_whitespace():
|
|
assert normalize_pci(" 0001:00:00.0 ") == (1, 0, 0, 0)
|
|
|
|
|
|
def test_normalize_pci_strips_0x_prefix():
|
|
assert normalize_pci("0x0001:00:00.0") == (1, 0, 0, 0)
|
|
|
|
|
|
def test_normalize_pci_missing_function_raises():
|
|
with pytest.raises(ValueError, match="missing function suffix"):
|
|
normalize_pci("0001:00:00")
|
|
|
|
|
|
def test_normalize_pci_invalid_function_char_raises():
|
|
with pytest.raises(ValueError, match="invalid PCI function"):
|
|
normalize_pci("0001:00:00.z")
|
|
|
|
|
|
def test_normalize_pci_too_many_segments_raises():
|
|
with pytest.raises(ValueError, match="invalid PCI BDF"):
|
|
normalize_pci("a:b:c:d.0")
|
|
|
|
|
|
def test_normalize_pci_bus_out_of_range_raises():
|
|
with pytest.raises(ValueError, match="out of range"):
|
|
normalize_pci("0000:1ff:00.0")
|
|
|
|
|
|
def test_normalize_pci_device_out_of_range_raises():
|
|
with pytest.raises(ValueError, match="out of range"):
|
|
normalize_pci("0000:00:20.0")
|
|
|
|
|
|
def test_normalize_pci_empty_string_raises():
|
|
with pytest.raises(ValueError):
|
|
normalize_pci("")
|
|
|
|
|
|
@pytest.mark.parametrize("pcp,pp", [(1, 1), (2, 1), (2, 2)])
|
|
def test_set_worker_net_device_uses_pcp_aware_gpu_index(monkeypatch, pcp, pp):
|
|
"""NIC lookup uses the same DP-offset GPU index as device selection."""
|
|
monkeypatch.setenv("VLLM_GPU_NIC_PCIE_MAPPING", "unused")
|
|
monkeypatch.setenv("VLLM_NIC_SELECTION_VARS", "unused")
|
|
seen: list[int] = []
|
|
monkeypatch.setattr(vllm_net_devices, "set_worker_gpu_nic_mapping", seen.append)
|
|
parallel_config = ParallelConfig(
|
|
tensor_parallel_size=2,
|
|
prefill_context_parallel_size=pcp,
|
|
pipeline_parallel_size=pp,
|
|
data_parallel_size=2,
|
|
data_parallel_rank_local=1,
|
|
distributed_executor_backend="mp",
|
|
)
|
|
vllm_net_devices.set_worker_net_device(
|
|
1, SimpleNamespace(parallel_config=parallel_config)
|
|
)
|
|
assert seen == [1 + 2 * pcp * pp]
|