[Distributed] Make it clear that % should not be in tensor dict keys. (#5927)

Signed-off-by: Xiaowei Jiang <xwjiang2010@gmail.com>
This commit is contained in:
xwjiang2010
2024-06-28 08:20:22 -07:00
committed by GitHub
parent 3b752a6555
commit b90d8cd832
2 changed files with 12 additions and 1 deletions

View File

@@ -1,5 +1,6 @@
from typing import Any, Dict
import pytest
import torch
from vllm.distributed.parallel_state import (_split_tensor_dict,
@@ -24,6 +25,14 @@ def test_split_tensor_dict():
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"),
@@ -31,7 +40,6 @@ def test_update_nested_dict():
("key8", "value5")]
res: Dict[str, Any] = {}
# Update the nested dictionary with each flattened key-value pair
for flat_key, value in flattened_keys_values:
_update_nested_dict(res, flat_key, value)
assert res == {