Spaces:
Running
Running
| """The bounded model cache: lazy loading, least recently used out, budget kept.""" | |
| import sys | |
| from pathlib import Path | |
| sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) | |
| from webapp.model_cache import ModelCache # noqa: E402 | |
| SIZES = {"a": 400.0, "b": 500.0, "c": 1000.0, "big": 5000.0} | |
| def _cache(budget): | |
| loads = [] | |
| cache = ModelCache(lambda path: loads.append(path) or {"path": path}, budget=budget, | |
| estimate=lambda path: SIZES[path]) | |
| return cache, loads | |
| def test_loads_on_first_use_only(): | |
| cache, loads = _cache(2000) | |
| assert cache.loaded_ids() == [] | |
| cache.get("a", "a") | |
| cache.get("a", "a") | |
| assert loads == ["a"] and cache.loaded_ids() == ["a"] | |
| def test_least_recently_used_goes_first(): | |
| cache, loads = _cache(1500) | |
| cache.get("a", "a") | |
| cache.get("b", "b") | |
| cache.get("a", "a") # a is now the most recent | |
| cache.get("c", "c") # 400 + 500 + 1000 > 1500: b goes, then a still fits? 400 + 1000 <= 1500 | |
| assert cache.loaded_ids() == ["a", "c"] | |
| assert cache.used_mb() <= 1500 | |
| def test_a_model_larger_than_the_budget_is_still_loaded_alone(): | |
| cache, loads = _cache(1500) | |
| cache.get("a", "a") | |
| cache.get("big", "big") | |
| assert cache.loaded_ids() == ["big"] | |
| def test_budget_from_environment(monkeypatch): | |
| from webapp import model_cache | |
| monkeypatch.setenv(model_cache.BUDGET_ENV, "1234") | |
| assert ModelCache(lambda p: p).budget == 1234 | |
| monkeypatch.setenv(model_cache.BUDGET_ENV, "lots") | |
| assert ModelCache(lambda p: p).budget == model_cache.DEFAULT_BUDGET_MB | |