from langchain_core.messages import HumanMessage, SystemMessage from arena.local_model import ( CachedLocalChatModel, _bucket_cache_key, clear_local_model_cache, local_model_cache_size, _cached_pipeline, _local_generation_duration, _model_load_sources, _messages_to_prompt, ) class FakePipeline: tokenizer = None def __init__(self) -> None: self.calls = 0 def __call__(self, prompt, **kwargs): self.calls += 1 return [{"generated_text": f"ok:{prompt[-10:]}"}] def test_cached_pipeline_reuses_loaded_model(monkeypatch) -> None: clear_local_model_cache() loads = [] def fake_load(model_id: str): loads.append(model_id) return FakePipeline() monkeypatch.setattr("arena.local_model._load_pipeline", fake_load) first = _cached_pipeline(model_id="local-7b", cache_ttl_seconds=1800) second = _cached_pipeline(model_id="local-7b", cache_ttl_seconds=1800) assert first is second assert loads == ["local-7b"] assert local_model_cache_size() == 1 clear_local_model_cache() def test_cached_pipeline_refresh_clears_model(monkeypatch) -> None: clear_local_model_cache() monkeypatch.setattr("arena.local_model._load_pipeline", lambda model_id: FakePipeline()) _cached_pipeline(model_id="local-7b", cache_ttl_seconds=1800) clear_local_model_cache() assert local_model_cache_size() == 0 def test_cached_local_chat_model_generates_with_cached_pipeline(monkeypatch) -> None: clear_local_model_cache() monkeypatch.setattr("arena.local_model._load_pipeline", lambda model_id: FakePipeline()) model = CachedLocalChatModel(model_id="local-7b", max_new_tokens=16) result = model.invoke([HumanMessage(content="Build a tiny app")]) assert result.content.startswith("ok:") assert local_model_cache_size() == 1 clear_local_model_cache() def test_local_gpu_duration_uses_cold_then_warm_cache(monkeypatch) -> None: clear_local_model_cache() monkeypatch.setenv("LOCAL_MODEL_GPU_COLD_DURATION_SECONDS", "700") monkeypatch.setenv("LOCAL_MODEL_GPU_WARM_DURATION_SECONDS", "90") monkeypatch.setattr("arena.local_model._load_pipeline", lambda model_id: FakePipeline()) cold = _local_generation_duration( model_id="local-7b", messages=[], temperature=0.2, max_new_tokens=16, stop=None, cache_ttl_seconds=1800, ) _cached_pipeline(model_id="local-7b", cache_ttl_seconds=1800) warm = _local_generation_duration( model_id="local-7b", messages=[], temperature=0.2, max_new_tokens=16, stop=None, cache_ttl_seconds=1800, ) assert cold == 700 assert warm == 90 clear_local_model_cache() def test_local_gpu_duration_does_not_sync_bucket(monkeypatch) -> None: clear_local_model_cache() monkeypatch.setenv("LOCAL_MODEL_BUCKET_URI", "hf://buckets/user/model-cache/gemma") def fail_sync(*args, **kwargs): raise AssertionError("duration lookup should not sync bucket files") monkeypatch.setattr("huggingface_hub.sync_bucket", fail_sync) duration = _local_generation_duration( model_id="local-7b", messages=[], temperature=0.2, max_new_tokens=16, stop=None, cache_ttl_seconds=1800, ) assert duration > 0 def test_model_load_sources_use_bucket_gguf(monkeypatch, tmp_path) -> None: bucket_uri = "hf://buckets/user/model-cache/gemma" cache_dir = tmp_path / "bucket-cache" monkeypatch.setenv("LOCAL_MODEL_BUCKET_URI", bucket_uri) monkeypatch.setenv("LOCAL_MODEL_BUCKET_CACHE_DIR", str(cache_dir)) monkeypatch.setenv("LOCAL_MODEL_GGUF_FILE", "model.gguf") def fake_sync(source, target): assert source == bucket_uri target_path = tmp_path / "bucket-cache" / _bucket_cache_key(bucket_uri) assert str(target_path) == target target_path.mkdir(parents=True, exist_ok=True) (target_path / "model.gguf").write_text("fake", encoding="utf-8") monkeypatch.setattr("huggingface_hub.sync_bucket", fake_sync) sources = _model_load_sources("unsloth/gemma-4-12B-it-qat-GGUF") assert sources.model_source == "unsloth/gemma-4-12B-it-qat-GGUF" assert sources.tokenizer_source == "unsloth/gemma-4-12B-it-qat-GGUF" assert sources.gguf_file.endswith("model.gguf") assert sources.bucket_uri == bucket_uri def test_model_load_sources_use_bucket_transformers_dir(monkeypatch, tmp_path) -> None: bucket_uri = "hf://buckets/user/model-cache/safetensors" cache_dir = tmp_path / "bucket-cache" monkeypatch.setenv("LOCAL_MODEL_BUCKET_URI", bucket_uri) monkeypatch.setenv("LOCAL_MODEL_BUCKET_CACHE_DIR", str(cache_dir)) def fake_sync(source, target): assert source == bucket_uri target_path = tmp_path / "bucket-cache" / _bucket_cache_key(bucket_uri) assert str(target_path) == target target_path.mkdir(parents=True, exist_ok=True) (target_path / "config.json").write_text("{}", encoding="utf-8") (target_path / "tokenizer.json").write_text("{}", encoding="utf-8") monkeypatch.setattr("huggingface_hub.sync_bucket", fake_sync) sources = _model_load_sources("local-bucket-model") assert sources.model_source == str(cache_dir / _bucket_cache_key(bucket_uri)) assert sources.tokenizer_source == str(cache_dir / _bucket_cache_key(bucket_uri)) assert sources.gguf_file == "" def test_model_load_sources_prefer_mounted_model_dir(monkeypatch, tmp_path) -> None: mounted_dir = tmp_path / "models" mounted_dir.mkdir() (mounted_dir / "model.gguf").write_text("fake", encoding="utf-8") monkeypatch.setenv("LOCAL_MODEL_BUCKET_URI", "hf://buckets/user/model-cache") monkeypatch.setenv("LOCAL_MODEL_MOUNT_DIR", str(mounted_dir)) monkeypatch.setenv("LOCAL_MODEL_GGUF_FILE", "model.gguf") def fail_sync(*args, **kwargs): raise AssertionError("mounted model dir should not sync bucket files") monkeypatch.setattr("huggingface_hub.sync_bucket", fail_sync) sources = _model_load_sources("unsloth/gemma-4-12B-it-qat-GGUF") assert sources.model_source == "unsloth/gemma-4-12B-it-qat-GGUF" assert sources.tokenizer_source == "unsloth/gemma-4-12B-it-qat-GGUF" assert sources.gguf_file == str(mounted_dir / "model.gguf") def test_messages_to_prompt_falls_back_without_chat_template() -> None: prompt = _messages_to_prompt( [ SystemMessage(content="You are concise."), HumanMessage(content="Build it."), ], tokenizer=None, ) assert "system: You are concise." in prompt assert "user: Build it." in prompt assert prompt.endswith("assistant:")