Merge 4e679e904f5ffee7b410eecb5340f50c0e8a97b6 into 553741557a57df44a46fb3ab72b6aafeaf945337

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]

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)