58 lines
1.6 KiB
Python
58 lines
1.6 KiB
Python
from typing import Any, Dict
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm.distributed.parallel_state import (_split_tensor_dict,
|
|
_update_nested_dict)
|
|
|
|
|
|
def test_split_tensor_dict():
|
|
test_dict = {
|
|
"key_a": "a",
|
|
"key_b": torch.arange(8, dtype=torch.float32),
|
|
"key_c": {
|
|
"key_1": torch.arange(5, dtype=torch.float32),
|
|
"key_2": torch.tensor([], dtype=torch.float32),
|
|
"key_3": 123,
|
|
},
|
|
"key_d": {},
|
|
}
|
|
metadata_list, tensor_list = _split_tensor_dict(test_dict)
|
|
assert len(metadata_list) == 6
|
|
assert torch.allclose(tensor_list[0], test_dict["key_b"])
|
|
assert torch.allclose(tensor_list[1], test_dict["key_c"]["key_1"])
|
|
assert torch.allclose(tensor_list[2], test_dict["key_c"]["key_2"])
|
|
|
|
|
|
def test_split_tensor_dict_invalid_key():
|
|
test_dict = {
|
|
"a%b": "a",
|
|
}
|
|
with pytest.raises(AssertionError):
|
|
_split_tensor_dict(test_dict)
|
|
|
|
|
|
def test_update_nested_dict():
|
|
flattened_keys_values = [("key1%key2%key3", "value1"),
|
|
("key1%key2%key4", "value2"),
|
|
("key1%key5", "value3"), ("key6%key7", "value4"),
|
|
("key8", "value5")]
|
|
res: Dict[str, Any] = {}
|
|
|
|
for flat_key, value in flattened_keys_values:
|
|
_update_nested_dict(res, flat_key, value)
|
|
assert res == {
|
|
"key1": {
|
|
"key2": {
|
|
"key3": "value1",
|
|
"key4": "value2"
|
|
},
|
|
"key5": "value3"
|
|
},
|
|
"key6": {
|
|
"key7": "value4"
|
|
},
|
|
"key8": "value5"
|
|
}
|