mirror of
https://github.com/agent0ai/agent-zero.git
synced 2026-10-14 03:38:53 +00:00
bea0c16d50
Run initialization end hooks while embedding preload continues, then invoke startup completion to perform the existing bounded wait and start MCP/jobs before HTTP. Preserve initialization and preload extension skip controls, best-effort failures, and the ten-second additional wait limit. Five balanced cached-model boots per version reduced median HTTP readiness from 17.063s to 15.907s (1.156s, 6.8%) and boot-to-first-token from 19.075s to 18.051s (1.024s, 5.4%). Useful overlap increased from 2.390s to 6.059s. Post-Send latency increased by 0.146s in this sample; local preparation increased by 0.025s. Models were warm before HTTP in all five candidate runs. Validation: 24 focused framework tests passed, including actual ten-second construction and encoding stalls, deterministic end-hook/warmup overlap, and extension skip/error paths. The clean live restart returned HTTP 200 and completed preload before HTTP. Temporary probes and benchmark chats were removed.
564 lines
20 KiB
Python
564 lines
20 KiB
Python
from concurrent.futures import ThreadPoolExecutor
|
|
import os
|
|
from pathlib import Path
|
|
import subprocess
|
|
import sys
|
|
import threading
|
|
import time
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
import models
|
|
|
|
|
|
def test_framework_imports_do_not_load_embedding_or_speech_runtimes():
|
|
result = subprocess.run(
|
|
[sys.executable, "-c", """
|
|
import sys
|
|
import models
|
|
import preload
|
|
for name in (
|
|
'sentence_transformers', 'torch', 'whisper',
|
|
'plugins._kokoro_tts.helpers.runtime',
|
|
'plugins._whisper_stt.helpers.runtime',
|
|
):
|
|
assert name not in sys.modules, name
|
|
"""],
|
|
cwd=Path(__file__).resolve().parents[1],
|
|
env={**os.environ, "LITELLM_LOCAL_MODEL_COST_MAP": "True"},
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=60,
|
|
)
|
|
assert result.returncode == 0, result.stderr
|
|
|
|
|
|
def _clear_local_embedding_models():
|
|
with models._LOCAL_EMBEDDING_MODELS_LOCK:
|
|
models._LOCAL_EMBEDDING_MODELS.clear()
|
|
|
|
|
|
@pytest.fixture
|
|
def startup_runtime(monkeypatch):
|
|
import run_ui
|
|
|
|
monkeypatch.setattr(run_ui, "configure_process_environment", lambda: None)
|
|
monkeypatch.setattr(run_ui, "run_migration_checks", lambda: None)
|
|
monkeypatch.setattr(run_ui, "prepare_web_runtime", lambda: (None, "localhost", 0))
|
|
monkeypatch.setattr(run_ui, "start_web_server", lambda *_args: None)
|
|
return run_ui
|
|
|
|
|
|
@pytest.mark.parametrize("skip_init", [False, True])
|
|
def test_runtime_preload_overlaps_chat_restoration_within_extension_hooks(monkeypatch, skip_init, startup_runtime):
|
|
import run_ui
|
|
|
|
calls = []
|
|
|
|
def extension_hook(point, agent, data):
|
|
if not point.startswith("_functions/run_ui/init_a0/"):
|
|
return
|
|
calls.append(point.rsplit("/", 1)[-1])
|
|
if skip_init and point.endswith("/start"):
|
|
data["result"] = None
|
|
|
|
monkeypatch.setattr(run_ui.extension, "call_extensions_sync", extension_hook)
|
|
def preload():
|
|
calls.append("preload")
|
|
|
|
def wait(timeout):
|
|
assert timeout == 10
|
|
calls.append("preload_ready")
|
|
|
|
return SimpleNamespace(result_sync=wait)
|
|
|
|
monkeypatch.setattr(run_ui.initialize, "initialize_preload", preload)
|
|
monkeypatch.setattr(
|
|
run_ui.initialize, "initialize_chats",
|
|
lambda: SimpleNamespace(result_sync=lambda: calls.append("chats")),
|
|
)
|
|
monkeypatch.setattr(run_ui.initialize, "initialize_mcp", lambda: calls.append("mcp"))
|
|
monkeypatch.setattr(run_ui.initialize, "initialize_job_loop", lambda: calls.append("jobs"))
|
|
|
|
run_ui.run()
|
|
|
|
assert calls == (["start", "end"] if skip_init else [
|
|
"start", "preload", "chats", "end", "preload_ready", "mcp", "jobs",
|
|
])
|
|
|
|
|
|
@pytest.mark.parametrize("error", [None, TimeoutError(), RuntimeError("preload failed")])
|
|
def test_runtime_preload_wait_remains_best_effort(monkeypatch, error, startup_runtime):
|
|
import run_ui
|
|
|
|
calls = []
|
|
errors = []
|
|
|
|
def wait(timeout):
|
|
assert timeout == 10
|
|
raise error
|
|
|
|
monkeypatch.setattr(run_ui.extension, "call_extensions_sync", lambda *_args, **_kwargs: None)
|
|
monkeypatch.setattr(
|
|
run_ui.initialize, "initialize_preload",
|
|
lambda: SimpleNamespace(result_sync=wait) if error is not None else None,
|
|
)
|
|
monkeypatch.setattr(
|
|
run_ui.initialize, "initialize_chats",
|
|
lambda: SimpleNamespace(result_sync=lambda: calls.append("chats")),
|
|
)
|
|
monkeypatch.setattr(run_ui.initialize, "initialize_mcp", lambda: calls.append("mcp"))
|
|
monkeypatch.setattr(run_ui.initialize, "initialize_job_loop", lambda: calls.append("jobs"))
|
|
monkeypatch.setattr(run_ui, "PrintStyle", lambda: SimpleNamespace(print=lambda _: None, error=errors.append))
|
|
|
|
run_ui.run()
|
|
|
|
assert calls == ["chats", "mcp", "jobs"]
|
|
assert errors == (["Error waiting for preload: preload failed"] if isinstance(error, RuntimeError) else [])
|
|
|
|
|
|
@pytest.mark.parametrize("blocked_stage", ["construction", "encoding"])
|
|
def test_startup_continues_after_ten_seconds_with_model_still_loading(monkeypatch, blocked_stage, record_property, startup_runtime):
|
|
import preload
|
|
import run_ui
|
|
from helpers import persist_chat
|
|
from helpers.defer import EventLoopThread
|
|
from plugins._model_config.helpers import model_config
|
|
|
|
release = threading.Event()
|
|
blocked = threading.Event()
|
|
chats_loaded = threading.Event()
|
|
warmed = threading.Event()
|
|
preload_done = threading.Event()
|
|
later_startup = []
|
|
errors = []
|
|
constructions = []
|
|
config = models.ModelConfig(models.ModelType.EMBEDDING, "huggingface", "sentence-transformers/timeout-test")
|
|
|
|
def extension_hook(point, agent, data):
|
|
if point.endswith("/initialize_preload/end"):
|
|
data["result"].add_done_callback(lambda _: preload_done.set())
|
|
|
|
def hold():
|
|
blocked.set()
|
|
assert release.wait(20), "Test did not release the model operation"
|
|
|
|
class HeldModel:
|
|
def __init__(self, name, **kwargs):
|
|
constructions.append(name)
|
|
if blocked_stage == "construction":
|
|
hold()
|
|
|
|
def encode(self, inputs, **kwargs):
|
|
if blocked_stage == "encoding":
|
|
hold()
|
|
warmed.set()
|
|
return [[1.0] for _ in inputs]
|
|
|
|
monkeypatch.setitem(sys.modules, "sentence_transformers", SimpleNamespace(SentenceTransformer=HeldModel))
|
|
monkeypatch.setattr(model_config, "get_embedding_model_config_object", lambda: config)
|
|
monkeypatch.setattr(models, "get_provider_config", lambda *_args: None)
|
|
monkeypatch.setattr(models, "get_api_key_raw", lambda *_args: "None")
|
|
monkeypatch.setattr(run_ui.extension, "call_extensions_sync", extension_hook)
|
|
monkeypatch.setattr(persist_chat, "load_tmp_chats", chats_loaded.set)
|
|
monkeypatch.setattr(run_ui.initialize, "initialize_mcp", lambda: later_startup.append("mcp"))
|
|
monkeypatch.setattr(run_ui.initialize, "initialize_job_loop", lambda: later_startup.append("jobs"))
|
|
monkeypatch.setattr(preload, "PrintStyle", lambda: SimpleNamespace(print=lambda _: None, error=errors.append))
|
|
monkeypatch.setattr(run_ui, "PrintStyle", lambda: SimpleNamespace(print=lambda _: None, error=errors.append))
|
|
_clear_local_embedding_models()
|
|
|
|
# Release blocked work even if a regression removes the startup timeout.
|
|
with ThreadPoolExecutor(max_workers=1) as pool:
|
|
started = time.monotonic()
|
|
startup = pool.submit(run_ui.run)
|
|
try:
|
|
assert blocked.wait(3)
|
|
assert chats_loaded.wait(3)
|
|
startup.result(timeout=12)
|
|
elapsed = time.monotonic() - started
|
|
record_property("startup_wait_seconds", round(elapsed, 3))
|
|
assert 9.5 <= elapsed < 13, elapsed
|
|
assert later_startup == ["mcp", "jobs"]
|
|
assert not warmed.is_set()
|
|
assert not release.is_set()
|
|
|
|
release.set()
|
|
assert preload_done.wait(3)
|
|
assert warmed.is_set()
|
|
runtime_model = models.get_embedding_model(config.provider, config.name, model_config=config)
|
|
assert runtime_model.embed_query("first message") == [1.0]
|
|
assert constructions == ["timeout-test"]
|
|
assert not errors
|
|
finally:
|
|
release.set()
|
|
startup.result(timeout=12)
|
|
assert preload_done.wait(3)
|
|
EventLoopThread().terminate()
|
|
_clear_local_embedding_models()
|
|
|
|
|
|
@pytest.mark.parametrize("scenario", ["cached", "missing", "broken", "remote", "busy_restoration"])
|
|
def test_real_preload_startup_handles_ready_or_unavailable_models(monkeypatch, scenario, record_property, startup_runtime):
|
|
import preload
|
|
import run_ui
|
|
from helpers import persist_chat
|
|
from helpers.defer import EventLoopThread
|
|
from plugins._model_config.helpers import model_config
|
|
|
|
chats_started = threading.Event()
|
|
release_chats = threading.Event()
|
|
model_constructed = threading.Event()
|
|
preload_done = threading.Event()
|
|
constructions = []
|
|
errors = []
|
|
calls = []
|
|
config = models.ModelConfig(
|
|
models.ModelType.EMBEDDING,
|
|
"openai" if scenario == "remote" else "huggingface",
|
|
"sentence-transformers/startup-test",
|
|
)
|
|
|
|
class Model:
|
|
def __init__(self, name, **kwargs):
|
|
constructions.append(name)
|
|
if scenario in {"missing", "broken"}:
|
|
raise OSError(scenario + " model")
|
|
model_constructed.set()
|
|
|
|
def encode(self, inputs, **kwargs):
|
|
return [[1.0] for _ in inputs]
|
|
|
|
def restore_chats():
|
|
chats_started.set()
|
|
if scenario == "busy_restoration":
|
|
assert release_chats.wait(5)
|
|
|
|
def extension_hook(point, agent, data):
|
|
if point.endswith("/initialize_preload/end"):
|
|
data["result"].add_done_callback(lambda _: preload_done.set())
|
|
|
|
monkeypatch.setitem(sys.modules, "sentence_transformers", SimpleNamespace(SentenceTransformer=Model))
|
|
monkeypatch.setattr(models, "get_provider_config", lambda *_args: None)
|
|
monkeypatch.setattr(models, "get_api_key_raw", lambda *_args: "None")
|
|
monkeypatch.setattr(model_config, "get_embedding_model_config_object", lambda: config)
|
|
monkeypatch.setattr(run_ui.extension, "call_extensions_sync", extension_hook)
|
|
monkeypatch.setattr(persist_chat, "load_tmp_chats", restore_chats)
|
|
monkeypatch.setattr(run_ui.initialize, "initialize_mcp", lambda: calls.append("mcp"))
|
|
monkeypatch.setattr(run_ui.initialize, "initialize_job_loop", lambda: calls.append("jobs"))
|
|
monkeypatch.setattr(preload, "PrintStyle", lambda: SimpleNamespace(print=lambda _: None, error=errors.append))
|
|
monkeypatch.setattr(run_ui, "PrintStyle", lambda: SimpleNamespace(print=lambda _: None, error=errors.append))
|
|
_clear_local_embedding_models()
|
|
if scenario == "cached":
|
|
models.get_embedding_model(config.provider, config.name)
|
|
|
|
with ThreadPoolExecutor(max_workers=1) as pool:
|
|
started = time.monotonic()
|
|
startup = pool.submit(run_ui.run)
|
|
try:
|
|
assert chats_started.wait(2)
|
|
if scenario == "busy_restoration":
|
|
assert model_constructed.wait(2)
|
|
assert not startup.done()
|
|
release_chats.set()
|
|
startup.result(timeout=2)
|
|
elapsed = time.monotonic() - started
|
|
record_property("startup_wait_seconds", round(elapsed, 3))
|
|
assert elapsed < 2, elapsed
|
|
assert preload_done.wait(2)
|
|
assert calls == ["mcp", "jobs"]
|
|
assert len(constructions) == (0 if scenario == "remote" else 1)
|
|
assert errors == (["Error in preload: " + scenario + " model"] if scenario in {"missing", "broken"} else [])
|
|
finally:
|
|
release_chats.set()
|
|
startup.result(timeout=12)
|
|
assert preload_done.wait(2)
|
|
EventLoopThread().terminate()
|
|
_clear_local_embedding_models()
|
|
|
|
|
|
def test_local_embedding_preload_is_reused_with_runtime_model_config(monkeypatch):
|
|
created = []
|
|
|
|
class FakeSentenceTransformer:
|
|
def __init__(self, model, **kwargs):
|
|
created.append((model, kwargs))
|
|
|
|
monkeypatch.setitem(
|
|
sys.modules, "sentence_transformers",
|
|
SimpleNamespace(SentenceTransformer=FakeSentenceTransformer),
|
|
)
|
|
_clear_local_embedding_models()
|
|
|
|
try:
|
|
preload = models.LocalSentenceTransformerWrapper(
|
|
"huggingface",
|
|
"sentence-transformers/example",
|
|
device="cpu",
|
|
model_kwargs={"revision": "stable", "trust_remote_code": False},
|
|
)
|
|
runtime_config = SimpleNamespace(name="runtime")
|
|
runtime = models.LocalSentenceTransformerWrapper(
|
|
"huggingface",
|
|
"sentence-transformers/example",
|
|
model_config=runtime_config,
|
|
model_kwargs={"trust_remote_code": False, "revision": "stable"},
|
|
device="cpu",
|
|
)
|
|
|
|
assert runtime.model is preload.model
|
|
assert runtime.a0_model_conf is runtime_config
|
|
assert created == [
|
|
(
|
|
"example",
|
|
{
|
|
"device": "cpu",
|
|
"model_kwargs": {
|
|
"revision": "stable",
|
|
"trust_remote_code": False,
|
|
},
|
|
},
|
|
)
|
|
]
|
|
finally:
|
|
_clear_local_embedding_models()
|
|
|
|
|
|
def test_local_embedding_cache_tracks_effective_constructor_options(monkeypatch):
|
|
created = []
|
|
|
|
class FakeSentenceTransformer:
|
|
def __init__(self, model, **kwargs):
|
|
created.append((model, kwargs))
|
|
|
|
monkeypatch.setitem(
|
|
sys.modules, "sentence_transformers",
|
|
SimpleNamespace(SentenceTransformer=FakeSentenceTransformer),
|
|
)
|
|
_clear_local_embedding_models()
|
|
|
|
try:
|
|
first = models.LocalSentenceTransformerWrapper(
|
|
"huggingface", "sentence-transformers/example", device="cpu"
|
|
)
|
|
second = models.LocalSentenceTransformerWrapper(
|
|
"huggingface", "sentence-transformers/example", device="cuda"
|
|
)
|
|
|
|
assert second.model is not first.model
|
|
assert created == [
|
|
("example", {"device": "cpu"}),
|
|
("example", {"device": "cuda"}),
|
|
]
|
|
assert len(models._LOCAL_EMBEDDING_MODELS) == 1
|
|
finally:
|
|
_clear_local_embedding_models()
|
|
|
|
|
|
def test_concurrent_preload_and_runtime_share_one_model(monkeypatch):
|
|
created = []
|
|
construction_started = threading.Event()
|
|
release_construction = threading.Event()
|
|
|
|
class FakeSentenceTransformer:
|
|
def __init__(self, model, **kwargs):
|
|
created.append((model, kwargs))
|
|
construction_started.set()
|
|
assert release_construction.wait(timeout=2)
|
|
|
|
monkeypatch.setitem(
|
|
sys.modules, "sentence_transformers",
|
|
SimpleNamespace(SentenceTransformer=FakeSentenceTransformer),
|
|
)
|
|
_clear_local_embedding_models()
|
|
|
|
try:
|
|
with ThreadPoolExecutor(max_workers=2) as executor:
|
|
first = executor.submit(
|
|
models.LocalSentenceTransformerWrapper,
|
|
"huggingface",
|
|
"sentence-transformers/example",
|
|
)
|
|
assert construction_started.wait(timeout=2)
|
|
second = executor.submit(
|
|
models.LocalSentenceTransformerWrapper,
|
|
"huggingface",
|
|
"sentence-transformers/example",
|
|
)
|
|
release_construction.set()
|
|
|
|
assert second.result().model is first.result().model
|
|
|
|
assert created == [("example", {})]
|
|
finally:
|
|
release_construction.set()
|
|
_clear_local_embedding_models()
|
|
|
|
|
|
def test_failed_model_change_keeps_the_working_cached_model(monkeypatch):
|
|
created = []
|
|
|
|
class FakeSentenceTransformer:
|
|
def __init__(self, model, **kwargs):
|
|
created.append((model, kwargs))
|
|
if model == "broken":
|
|
raise RuntimeError("model unavailable")
|
|
|
|
monkeypatch.setitem(
|
|
sys.modules, "sentence_transformers",
|
|
SimpleNamespace(SentenceTransformer=FakeSentenceTransformer),
|
|
)
|
|
_clear_local_embedding_models()
|
|
|
|
try:
|
|
working = models.LocalSentenceTransformerWrapper(
|
|
"huggingface", "sentence-transformers/working"
|
|
)
|
|
with pytest.raises(RuntimeError, match="model unavailable"):
|
|
models.LocalSentenceTransformerWrapper(
|
|
"huggingface", "sentence-transformers/broken"
|
|
)
|
|
reused = models.LocalSentenceTransformerWrapper(
|
|
"huggingface", "sentence-transformers/working"
|
|
)
|
|
|
|
assert reused.model is working.model
|
|
assert created == [("working", {}), ("broken", {})]
|
|
finally:
|
|
_clear_local_embedding_models()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_preload_uses_the_runtime_embedding_configuration(monkeypatch):
|
|
import preload
|
|
from plugins._model_config.helpers import model_config
|
|
|
|
config = SimpleNamespace(
|
|
provider="huggingface",
|
|
name="sentence-transformers/example",
|
|
build_kwargs=lambda: {"device": "cpu"},
|
|
)
|
|
calls = []
|
|
embedded = []
|
|
loop_thread = threading.get_ident()
|
|
|
|
class FakeEmbeddings:
|
|
async def aembed_query(self, text):
|
|
embedded.append(text)
|
|
|
|
def get_embedding_model(provider, name, **kwargs):
|
|
assert threading.get_ident() != loop_thread
|
|
calls.append((provider, name, kwargs))
|
|
return FakeEmbeddings()
|
|
|
|
monkeypatch.setattr(
|
|
model_config, "get_embedding_model_config_object", lambda: config
|
|
)
|
|
monkeypatch.setattr(preload.models, "get_embedding_model", get_embedding_model)
|
|
|
|
await preload.preload()
|
|
|
|
assert calls == [
|
|
(
|
|
"huggingface",
|
|
"sentence-transformers/example",
|
|
{"model_config": config, "device": "cpu"},
|
|
)
|
|
]
|
|
assert embedded == ["test"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("provider", ["openai", "huggingface"])
|
|
async def test_preload_skips_remote_models_and_contains_local_failures(monkeypatch, provider):
|
|
import preload
|
|
from plugins._model_config.helpers import model_config
|
|
|
|
config = SimpleNamespace(provider=provider, name="example", build_kwargs=lambda: {})
|
|
calls = []
|
|
errors = []
|
|
|
|
def fail_embedding(*args, **kwargs):
|
|
calls.append(args)
|
|
raise RuntimeError("model unavailable")
|
|
|
|
monkeypatch.setattr(model_config, "get_embedding_model_config_object", lambda: config)
|
|
monkeypatch.setattr(models, "get_embedding_model", fail_embedding)
|
|
monkeypatch.setattr(
|
|
preload, "PrintStyle",
|
|
lambda: SimpleNamespace(print=lambda _: None, error=errors.append),
|
|
)
|
|
|
|
await preload.preload()
|
|
|
|
if provider == "huggingface":
|
|
assert calls == [("huggingface", "example")]
|
|
assert errors == ["Error in preload: model unavailable"]
|
|
else:
|
|
assert calls == []
|
|
assert errors == []
|
|
|
|
|
|
def test_watchdog_setup_overlaps_real_warmup_before_mcp_and_http(monkeypatch, startup_runtime):
|
|
import preload
|
|
import run_ui
|
|
from helpers import persist_chat
|
|
from helpers.defer import EventLoopThread
|
|
from plugins._model_config.helpers import model_config
|
|
|
|
model_started = threading.Event()
|
|
release_model = threading.Event()
|
|
preload_done = threading.Event()
|
|
calls = []
|
|
config = models.ModelConfig(
|
|
models.ModelType.EMBEDDING, "huggingface", "sentence-transformers/overlap-test"
|
|
)
|
|
|
|
class Model:
|
|
def __init__(self, *_args, **_kwargs):
|
|
model_started.set()
|
|
assert release_model.wait(5)
|
|
|
|
def encode(self, inputs, **_kwargs):
|
|
calls.append("model_warmed")
|
|
return [[1.0] for _ in inputs]
|
|
|
|
def extension_hook(point, agent, data):
|
|
if point.endswith("/initialize_preload/end"):
|
|
data["result"].add_done_callback(lambda _: preload_done.set())
|
|
if point.endswith("/init_a0/end"):
|
|
assert model_started.wait(2)
|
|
assert not preload_done.is_set()
|
|
calls.append("watchdogs_start")
|
|
release_model.set()
|
|
assert preload_done.wait(2)
|
|
calls.append("watchdogs_end")
|
|
|
|
monkeypatch.setitem(sys.modules, "sentence_transformers", SimpleNamespace(SentenceTransformer=Model))
|
|
monkeypatch.setattr(model_config, "get_embedding_model_config_object", lambda: config)
|
|
monkeypatch.setattr(models, "get_provider_config", lambda *_args: None)
|
|
monkeypatch.setattr(models, "get_api_key_raw", lambda *_args: "None")
|
|
monkeypatch.setattr(run_ui.extension, "call_extensions_sync", extension_hook)
|
|
monkeypatch.setattr(persist_chat, "load_tmp_chats", lambda: calls.append("chats"))
|
|
monkeypatch.setattr(run_ui.initialize, "initialize_mcp", lambda: calls.append("mcp"))
|
|
monkeypatch.setattr(run_ui.initialize, "initialize_job_loop", lambda: calls.append("jobs"))
|
|
monkeypatch.setattr(run_ui, "start_web_server", lambda *_args: calls.append("http"))
|
|
_clear_local_embedding_models()
|
|
try:
|
|
run_ui.run()
|
|
assert calls == ["chats", "watchdogs_start", "model_warmed", "watchdogs_end", "mcp", "jobs", "http"]
|
|
finally:
|
|
release_model.set()
|
|
assert preload_done.wait(2)
|
|
EventLoopThread().terminate()
|
|
_clear_local_embedding_models()
|
|
|
|
|
|
def test_runtime_accepts_extension_skip_result(monkeypatch, startup_runtime):
|
|
import run_ui
|
|
|
|
calls = []
|
|
monkeypatch.setattr(run_ui, "init_a0", lambda: True)
|
|
monkeypatch.setattr(run_ui, "start_web_server", lambda *_args: calls.append("http"))
|
|
run_ui.run()
|
|
assert calls == ["http"]
|