mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-31 00:36:20 +08:00
37 lines
1.4 KiB
Python
37 lines
1.4 KiB
Python
import json
|
|
import time
|
|
from collections.abc import Iterator
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
import torch.distributed as dist
|
|
|
|
from deepinfer import envs
|
|
from deepinfer.config import ParallelConfig
|
|
from deepinfer.worker.decode_perf_cache import DecodePerfCache
|
|
|
|
pytestmark = [pytest.mark.ci, pytest.mark.ci_npu]
|
|
|
|
|
|
@pytest.fixture
|
|
def cache(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Iterator[DecodePerfCache]:
|
|
monkeypatch.setattr(envs, "DEEPINFER_DECODE_PERF_CACHE_DIR", str(tmp_path / "cache"))
|
|
dist.init_process_group("gloo", init_method=(tmp_path / "rendezvous").as_uri(), rank=0, world_size=1)
|
|
try:
|
|
config = ParallelConfig(1, 1, 1, 1, 0, 0)
|
|
yield DecodePerfCache("/models/model", config, [1], dist.group.WORLD)
|
|
finally:
|
|
dist.destroy_process_group()
|
|
|
|
|
|
@pytest.mark.parametrize("batch_size", ["1e309", "-1e309", "Infinity", "-Infinity"])
|
|
@pytest.mark.parametrize("has_older_curve", [False, True])
|
|
def test_overflowing_batch_size_is_a_cache_miss(cache: DecodePerfCache, batch_size: str, has_older_curve: bool) -> None:
|
|
folder = Path(cache.cache_dir)
|
|
folder.mkdir()
|
|
now = int(time.time())
|
|
if has_older_curve:
|
|
(folder / f"{cache.hash_key}.{'a' * 32}.{now - 1}").write_text(json.dumps([[1, 2.5]]))
|
|
(folder / f"{cache.hash_key}.{'b' * 32}.{now}").write_text(f"[[{batch_size}, 3.0]]")
|
|
assert cache.load_decode_perf_cache() == ([(1, 2.5)] if has_older_curve else None)
|