File size: 6,779 Bytes
e140461 7d806c0 e140461 | 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 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 | 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:")
|