From 2afffe8edcb82b17ae1fe44e637b22182f6c3b84 Mon Sep 17 00:00:00 2001 From: Alessandro <155005371+3clyp50@users.noreply.github.com> Date: Sat, 10 Oct 2026 16:36:08 +0200 Subject: [PATCH] Overlap embedding warm-up with chat restoration Load sentence-transformers on the existing cache miss and construct the preloaded model off the shared background event loop. Start warm-up before chat restoration inside extensible init_a0, preserving start-hook control and sharing the loaded model with ordinary runtime requests. Remove unused speech preload imports and dead warm-up functions. Seed APT indexes only when absent, skipping the refresh on normal container boots. Keep the existing Python tunnel service and its provider lifecycle. Validation: 49 focused framework tests passed, including extension-hook short-circuit behavior, concurrent embedding-cache reuse, APT cache states, tunnel behavior and CSRF origins. Three cold Fable chat trials per version: median submit-to-first-token 9.588s -> 4.234s; local preparation 6.904s -> 1.197s; HTTP readiness 42.637s -> 41.139s; boot-to-first-token 52.093s -> 46.657s. Provider timing varies; these are local repeated measurements. The live baseline already omitted APT refresh, so it receives no timing credit. --- AGENTS.md | 1 + docker/run/AGENTS.md | 1 + docker/run/fs/exe/initialize.sh | 9 ++- models.py | 11 ++- preload.py | 59 +++----------- run_ui.py | 2 +- tests/test_docker_initialize_limits.py | 36 +++++++-- tests/test_embedding_model_preload.py | 106 ++++++++++++++++++++++++- 8 files changed, 164 insertions(+), 61 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 48dea25e4..65fe58df2 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -22,6 +22,7 @@ - Native function calls use canonical call content for history and combine public response text with calls for repeat comparison. The canonical history content passes through the normal history template/masking hook; provider output metadata remains intact. - `initialize.py` owns framework initialization. - `models.py` owns model-provider configuration and LiteLLM integration. `get_api_key_raw()` reads the stored key or key list for settings/editing; `get_api_key()` selects a runtime key and runs provider extension hooks. +- Local sentence-transformers load only on an embedding cache miss, not when importing `models`. `run_ui.init_a0` starts the deferred preload before chat restoration, overlapping model loading while preserving the initialization extension hooks. `preload.py` constructs the configured local embedding model in a worker thread; it must not import unused speech runtimes or block the shared background event loop. - `models._merge_provider_defaults` merges YAML-declared `extra_body` defaults with legacy top-level values for those declared fields and explicit nested overrides, in that order; unrelated kwargs and original-provider API-key selection remain intact. - `models.apply_rate_limiter` skips token counting and returns no limiter when all request, input, and output limits are unset; any configured limit retains normal accounting. - `requirements.txt` includes HTTPX's SOCKS extra so framework clients support SOCKS proxy environment variables before plugin setup. diff --git a/docker/run/AGENTS.md b/docker/run/AGENTS.md index d1ae9b86b..69447feae 100644 --- a/docker/run/AGENTS.md +++ b/docker/run/AGENTS.md @@ -22,6 +22,7 @@ - For v2.13 builds, retain the Python 3.13-compatible ATK, LibreOffice/UNO, and Xpra pins in `fs/ins/install_additional.sh`. Resolve them from the signed Debian `forky` archive dated `20260624T000000Z` and Xpra `trixie`, using build-scoped APT options that preserve normal Kali sources. - Do not bake secrets, local `.env` values, or user data into the image. - Runtime startup must ensure `/a0/usr/uploads` exists before supervised services start. +- Runtime startup seeds APT package lists in the background only when no package index is cached (fresh images clear these at build time). Normal restarts skip the refresh; installation scripts and plugin hooks refresh metadata when installing packages. - Runtime startup raises the soft open-file limit toward `A0_NOFILE_LIMIT` (default `65535`) before supervisord starts, bounded by the container hard limit. - Self-update user-data backups skip Time Travel shadow history under `usr/.time_travel/` and transient Desktop agent state. - Self-update rollback stashes use immutable Git object IDs for creation checks and restoration; resolve the matching reflog selector only when dropping that stash, preserving unrelated entries. Failed restoration retains the stash. diff --git a/docker/run/fs/exe/initialize.sh b/docker/run/fs/exe/initialize.sh index f4c4dcd25..b5a9dce75 100644 --- a/docker/run/fs/exe/initialize.sh +++ b/docker/run/fs/exe/initialize.sh @@ -44,6 +44,12 @@ raise_open_file_limit() { fi } +initialize_package_lists() { + if ! compgen -G "/var/lib/apt/lists/*_Packages*" > /dev/null; then + apt-get update > /dev/null 2>&1 & + fi +} + raise_open_file_limit # Copy all contents from persistent /per to root directory (/) without overwriting @@ -56,8 +62,7 @@ mkdir -p /a0/usr/uploads chmod 444 /root/.bashrc chmod 444 /root/.profile -# update package list to save time later -apt-get update > /dev/null 2>&1 & +initialize_package_lists # let supervisord handle the services exec /usr/bin/supervisord -c /etc/supervisor/conf.d/supervisord.conf diff --git a/models.py b/models.py index 9b3cf3ecd..53a9d1973 100644 --- a/models.py +++ b/models.py @@ -14,6 +14,7 @@ from typing import ( AsyncIterator, Tuple, TypedDict, + TYPE_CHECKING, ) from litellm import embedding @@ -43,9 +44,11 @@ from langchain_core.messages import ( SystemMessage, ) from langchain.embeddings.base import Embeddings -from sentence_transformers import SentenceTransformer from pydantic import ConfigDict +if TYPE_CHECKING: + from sentence_transformers import SentenceTransformer + DEFAULT_LITELLM_GLOBAL_KWARGS: dict[str, Any] = { "drop_params": True, @@ -832,17 +835,19 @@ class LiteLLMEmbeddingWrapper(Embeddings): return self.embed([text])[0] -_LOCAL_EMBEDDING_MODELS: dict[tuple[str, str], SentenceTransformer] = {} +_LOCAL_EMBEDDING_MODELS: dict[tuple[str, str], "SentenceTransformer"] = {} _LOCAL_EMBEDDING_MODELS_LOCK = threading.Lock() def _get_local_embedding_model( model: str, kwargs: dict[str, Any] -) -> SentenceTransformer: +) -> "SentenceTransformer": key = (model, json.dumps(kwargs, sort_keys=True, default=repr)) with _LOCAL_EMBEDDING_MODELS_LOCK: cached = _LOCAL_EMBEDDING_MODELS.get(key) if cached is None: + from sentence_transformers import SentenceTransformer + cached = SentenceTransformer(model, **kwargs) _LOCAL_EMBEDDING_MODELS.clear() _LOCAL_EMBEDDING_MODELS[key] = cached diff --git a/preload.py b/preload.py index 914c41969..e06432ee5 100644 --- a/preload.py +++ b/preload.py @@ -2,62 +2,29 @@ import asyncio from helpers import runtime from helpers.print_style import PrintStyle import models -from plugins._kokoro_tts.helpers import runtime as kokoro_tts_runtime -from plugins._whisper_stt.helpers import runtime as whisper_stt_runtime async def preload(): try: - # preload whisper model - async def preload_whisper(): - if not whisper_stt_runtime.is_globally_enabled(): - return None - try: - config = whisper_stt_runtime.get_config() - return await whisper_stt_runtime.preload(str(config["model_size"])) - except Exception as e: - PrintStyle().error(f"Error in preload_whisper: {e}") + from plugins._model_config.helpers.model_config import get_embedding_model_config_object - # preload embedding model - async def preload_embedding(): - try: - from plugins._model_config.helpers.model_config import get_embedding_model_config_object - emb_cfg = get_embedding_model_config_object() - if emb_cfg.provider.lower() == "huggingface": - emb_mod = models.get_embedding_model( - emb_cfg.provider, - emb_cfg.name, - model_config=emb_cfg, - **emb_cfg.build_kwargs(), - ) - emb_txt = await emb_mod.aembed_query("test") - return emb_txt - except Exception as e: - PrintStyle().error(f"Error in preload_embedding: {e}") - - # preload kokoro tts model if enabled - async def preload_kokoro(): - if not kokoro_tts_runtime.is_globally_enabled(): - return None - try: - return await kokoro_tts_runtime.preload() - except Exception as e: - PrintStyle().error(f"Error in preload_kokoro: {e}") - - # async tasks to preload - tasks = [ - preload_embedding(), - # preload_whisper(), - # preload_kokoro() - ] - - await asyncio.gather(*tasks, return_exceptions=True) + emb_cfg = get_embedding_model_config_object() + if emb_cfg.provider.lower() == "huggingface": + # Model imports and construction must not block the shared background loop. + emb_mod = await asyncio.to_thread( + models.get_embedding_model, + emb_cfg.provider, + emb_cfg.name, + model_config=emb_cfg, + **emb_cfg.build_kwargs(), + ) + await emb_mod.aembed_query("test") PrintStyle().print("Preload completed.") except Exception as e: PrintStyle().error(f"Error in preload: {e}") -# preload transcription model +# preload the configured embedding model if __name__ == "__main__": PrintStyle().print("Running preload...") runtime.initialize() diff --git a/run_ui.py b/run_ui.py index 8346d1b47..b9360f5d3 100644 --- a/run_ui.py +++ b/run_ui.py @@ -73,12 +73,12 @@ def create_flush_callback(): @extension.extensible def init_a0(): + initialize.initialize_preload() init_chats = initialize.initialize_chats() init_chats.result_sync() initialize.initialize_mcp() initialize.initialize_job_loop() - initialize.initialize_preload() if __name__ == "__main__": diff --git a/tests/test_docker_initialize_limits.py b/tests/test_docker_initialize_limits.py index f4c149d6e..6fb83a6c0 100644 --- a/tests/test_docker_initialize_limits.py +++ b/tests/test_docker_initialize_limits.py @@ -10,10 +10,10 @@ REPO_ROOT = Path(__file__).resolve().parents[1] INITIALIZE_SCRIPT = REPO_ROOT / "docker" / "run" / "fs" / "exe" / "initialize.sh" -def _raise_limit_function() -> str: +def _script_function(name: str) -> str: text = INITIALIZE_SCRIPT.read_text(encoding="utf-8") - match = re.search(r"^raise_open_file_limit\(\) \{\n.*?^\}\n", text, re.M | re.S) - assert match, "initialize.sh must define raise_open_file_limit" + match = re.search(rf"^{name}\(\) \{{\n.*?^\}}\n", text, re.M | re.S) + assert match, f"initialize.sh must define {name}" return match.group(0) @@ -27,7 +27,7 @@ def _run_bash(script: str) -> subprocess.CompletedProcess[str]: def test_initialize_raises_soft_open_file_limit_to_requested_target(): - function = _raise_limit_function() + function = _script_function("raise_open_file_limit") result = _run_bash( f""" @@ -48,7 +48,7 @@ def test_initialize_caps_open_file_limit_at_hard_limit(): if hard != resource.RLIM_INFINITY and hard < 2048: pytest.skip("host hard open-file limit is too low for this regression test") - function = _raise_limit_function() + function = _script_function("raise_open_file_limit") result = _run_bash( f""" @@ -63,3 +63,29 @@ def test_initialize_caps_open_file_limit_at_hard_limit(): ) assert "Raised open file soft limit from 1024 to 2048" in result.stdout + + +@pytest.mark.parametrize("package_index", [None, "mirror_Packages", "mirror_Packages.lz4"]) +def test_initialize_updates_only_an_empty_package_cache(tmp_path, package_index): + lists = tmp_path / "lists" + lists.mkdir() + (lists / "partial").mkdir() + (lists / "lock").touch() + if package_index: + (lists / package_index).write_text("cached package index") + calls = tmp_path / "apt-calls" + function = _script_function("initialize_package_lists").replace( + "/var/lib/apt/lists", str(lists) + ) + + _run_bash(f""" + set -euo pipefail + apt-get() {{ echo "$*" >> "{calls}"; }} + {function} + initialize_package_lists + wait + """) + + assert (calls.read_text() if calls.exists() else "") == ( + "" if package_index else "update\n" + ) diff --git a/tests/test_embedding_model_preload.py b/tests/test_embedding_model_preload.py index 50f4634de..5071c0f33 100644 --- a/tests/test_embedding_model_preload.py +++ b/tests/test_embedding_model_preload.py @@ -1,4 +1,8 @@ from concurrent.futures import ThreadPoolExecutor +import os +from pathlib import Path +import subprocess +import sys import threading from types import SimpleNamespace @@ -7,11 +11,60 @@ 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.mark.parametrize("skip_init", [False, True]) +def test_runtime_preload_overlaps_chat_restoration_within_extension_hooks(monkeypatch, skip_init): + import run_ui + + calls = [] + + def extension_hook(point, agent, data): + 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) + monkeypatch.setattr(run_ui.initialize, "initialize_preload", lambda: calls.append("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.init_a0() + + assert calls == (["start", "end"] if skip_init else [ + "start", "preload", "chats", "mcp", "jobs", "end", + ]) + + def test_local_embedding_preload_is_reused_with_runtime_model_config(monkeypatch): created = [] @@ -19,7 +72,10 @@ def test_local_embedding_preload_is_reused_with_runtime_model_config(monkeypatch def __init__(self, model, **kwargs): created.append((model, kwargs)) - monkeypatch.setattr(models, "SentenceTransformer", FakeSentenceTransformer) + monkeypatch.setitem( + sys.modules, "sentence_transformers", + SimpleNamespace(SentenceTransformer=FakeSentenceTransformer), + ) _clear_local_embedding_models() try: @@ -63,7 +119,10 @@ def test_local_embedding_cache_tracks_effective_constructor_options(monkeypatch) def __init__(self, model, **kwargs): created.append((model, kwargs)) - monkeypatch.setattr(models, "SentenceTransformer", FakeSentenceTransformer) + monkeypatch.setitem( + sys.modules, "sentence_transformers", + SimpleNamespace(SentenceTransformer=FakeSentenceTransformer), + ) _clear_local_embedding_models() try: @@ -95,7 +154,10 @@ def test_concurrent_preload_and_runtime_share_one_model(monkeypatch): construction_started.set() assert release_construction.wait(timeout=2) - monkeypatch.setattr(models, "SentenceTransformer", FakeSentenceTransformer) + monkeypatch.setitem( + sys.modules, "sentence_transformers", + SimpleNamespace(SentenceTransformer=FakeSentenceTransformer), + ) _clear_local_embedding_models() try: @@ -130,7 +192,10 @@ def test_failed_model_change_keeps_the_working_cached_model(monkeypatch): if model == "broken": raise RuntimeError("model unavailable") - monkeypatch.setattr(models, "SentenceTransformer", FakeSentenceTransformer) + monkeypatch.setitem( + sys.modules, "sentence_transformers", + SimpleNamespace(SentenceTransformer=FakeSentenceTransformer), + ) _clear_local_embedding_models() try: @@ -163,12 +228,14 @@ async def test_preload_uses_the_runtime_embedding_configuration(monkeypatch): ) 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() @@ -187,3 +254,34 @@ async def test_preload_uses_the_runtime_embedding_configuration(monkeypatch): ) ] 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 == []