File size: 1,895 Bytes
a402b9b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | import sys
from pathlib import Path
import pytest
import torch
from sglang.srt.debug_utils.dump_loader import read_tokenizer_path
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=10, suite="default", nightly=True)
def _save_pt(
directory: Path, filename: str, *, value: torch.Tensor, meta: dict
) -> None:
torch.save({"value": value, "meta": meta}, directory / filename)
class TestReadTokenizerPath:
def test_finds_tokenizer_path(self, tmp_path: Path) -> None:
_save_pt(
tmp_path,
"name=x___step=0___rank=0___dump_index=0.pt",
value=torch.tensor([1.0]),
meta={"tokenizer_path": "/models/llama-3"},
)
result = read_tokenizer_path(tmp_path)
assert result == "/models/llama-3"
def test_returns_none_when_no_tokenizer_path(self, tmp_path: Path) -> None:
_save_pt(
tmp_path,
"name=x___step=0___rank=0___dump_index=0.pt",
value=torch.tensor([1.0]),
meta={},
)
result = read_tokenizer_path(tmp_path)
assert result is None
def test_returns_none_for_empty_directory(self, tmp_path: Path) -> None:
result = read_tokenizer_path(tmp_path)
assert result is None
def test_skips_files_without_tokenizer_path(self, tmp_path: Path) -> None:
_save_pt(
tmp_path,
"name=a___step=0___rank=0___dump_index=0.pt",
value=torch.tensor([1.0]),
meta={},
)
_save_pt(
tmp_path,
"name=b___step=0___rank=0___dump_index=1.pt",
value=torch.tensor([2.0]),
meta={"tokenizer_path": "/models/deepseek"},
)
result = read_tokenizer_path(tmp_path)
assert result == "/models/deepseek"
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))
|