Spaces:
Paused
Paused
File size: 2,026 Bytes
ec14ba6 | 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 | """Verify model reuse and eviction without torch imports or model downloads."""
from types import SimpleNamespace
import unittest
from unittest.mock import Mock, patch
from inference import InferenceEngine
class ModelCacheTests(unittest.TestCase):
def test_two_model_lru_reuses_recent_tasks_and_evicts_before_load(self):
engine = InferenceEngine(model_cache_size=2)
observed_sizes = []
def load(*args, **kwargs):
observed_sizes.append(len(engine._cache))
return SimpleNamespace(eval=lambda: None)
model_class = SimpleNamespace(from_pretrained=Mock(side_effect=load))
transformers = SimpleNamespace(AutoModel=model_class, AutoModelForSequenceClassification=model_class,
AutoModelForTokenClassification=model_class)
torch = SimpleNamespace(float32="float32", set_num_threads=lambda count: None)
tokenizers = {name: object() for name in ["domain", "pii", "guard"]}
with patch.dict("sys.modules", {"torch": torch, "transformers": transformers}):
for task in ["domain", "pii", "domain", "guard"]:
engine._ensure_model(task, tokenizers[task])
self.assertEqual(model_class.from_pretrained.call_count, 3)
self.assertEqual(observed_sizes, [0, 1, 1])
self.assertEqual(engine.status()["cached_tasks"], ["domain", "guard"])
self.assertIs(engine._get_tokenizer("domain"), tokenizers["domain"])
self.assertEqual(engine.resident_task, "guard")
def test_single_model_configuration_is_supported_and_bounds_are_enforced(self):
self.assertEqual(InferenceEngine(model_cache_size=1).model_cache_size, 1)
with patch.dict("os.environ", {"VELA_MODEL_CACHE_SIZE": "1"}):
self.assertEqual(InferenceEngine().model_cache_size, 1)
for invalid in [0, 3]:
with self.assertRaises(ValueError):
InferenceEngine(model_cache_size=invalid)
if __name__ == "__main__":
unittest.main()
|