Spaces:
Paused
Paused
| """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() | |