diff --git a/deepinfer/worker/decode_perf_cache.py b/deepinfer/worker/decode_perf_cache.py index 843470b..d661093 100644 --- a/deepinfer/worker/decode_perf_cache.py +++ b/deepinfer/worker/decode_perf_cache.py @@ -112,7 +112,7 @@ class DecodePerfCache: with open(path) as f: raw_results = json.load(f) results = [(int(batch_size), float(ms)) for batch_size, ms in raw_results] - except (OSError, TypeError, ValueError, json.JSONDecodeError) as e: + except (OSError, OverflowError, TypeError, ValueError, json.JSONDecodeError) as e: logger.warning(f"Failed to read decode perf cache file {path}: {e}") continue batch_sizes = [batch_size for batch_size, _ in results] diff --git a/tests/worker/test_decode_perf_cache_overflow.py b/tests/worker/test_decode_perf_cache_overflow.py new file mode 100644 index 0000000..99457d7 --- /dev/null +++ b/tests/worker/test_decode_perf_cache_overflow.py @@ -0,0 +1,36 @@ +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)