mirror of
https://github.com/onyx-dot-app/onyx.git
synced 2026-10-19 01:50:51 -05:00
perf(llm): initialize litellm once per process, not once per message (#14964)
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
DanelegendandClaude Opus 5 authored and GitHub committed
1 parent
d34e8dcd4b
commit
1d72e59f06
4 files changed
+250
-96
No files matched your search
@@ -838,10 +838,8 @@ def run_llm_loop(
|
||||
user_id=user_identity.user_id if user_identity else None,
|
||||
).model_dump(),
|
||||
):
|
||||
# Fix some LiteLLM issues,
|
||||
from onyx.llm.litellm_singleton.config import (
|
||||
initialize_litellm,
|
||||
) # Here for lazy load LiteLLM
|
||||
# Here for lazy load LiteLLM. initialize_litellm runs once per process.
|
||||
from onyx.llm.litellm_singleton.config import initialize_litellm
|
||||
|
||||
initialize_litellm()
|
||||
|
||||
|
||||
@@ -228,7 +228,7 @@ def run_deep_research_llm_loop(
|
||||
user_id=user_identity.user_id if user_identity else None,
|
||||
).model_dump(),
|
||||
):
|
||||
# Here for lazy load LiteLLM
|
||||
# Here for lazy load LiteLLM. initialize_litellm runs once per process.
|
||||
from onyx.llm.litellm_singleton.config import initialize_litellm
|
||||
|
||||
# An approximate limit. In extreme cases it may still fail but this should allow deep research
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import json
|
||||
import threading
|
||||
from pathlib import Path
|
||||
|
||||
import litellm
|
||||
@@ -28,94 +29,98 @@ def configure_litellm_settings() -> None:
|
||||
|
||||
# TODO: We might not need to register ollama_chat in addition to ollama but let's just do it for good measure for now.
|
||||
def register_ollama_models() -> None:
|
||||
litellm.register_model(
|
||||
model_cost={
|
||||
# GPT-OSS models
|
||||
"ollama_chat/gpt-oss:120b-cloud": {"supports_function_calling": True},
|
||||
"ollama_chat/gpt-oss:120b": {"supports_function_calling": True},
|
||||
"ollama_chat/gpt-oss:20b-cloud": {"supports_function_calling": True},
|
||||
"ollama_chat/gpt-oss:20b": {"supports_function_calling": True},
|
||||
"ollama/gpt-oss:120b-cloud": {"supports_function_calling": True},
|
||||
"ollama/gpt-oss:120b": {"supports_function_calling": True},
|
||||
"ollama/gpt-oss:20b-cloud": {"supports_function_calling": True},
|
||||
"ollama/gpt-oss:20b": {"supports_function_calling": True},
|
||||
# DeepSeek models
|
||||
"ollama_chat/deepseek-r1:latest": {"supports_function_calling": True},
|
||||
"ollama_chat/deepseek-r1:1.5b": {"supports_function_calling": True},
|
||||
"ollama_chat/deepseek-r1:7b": {"supports_function_calling": True},
|
||||
"ollama_chat/deepseek-r1:8b": {"supports_function_calling": True},
|
||||
"ollama_chat/deepseek-r1:14b": {"supports_function_calling": True},
|
||||
"ollama_chat/deepseek-r1:32b": {"supports_function_calling": True},
|
||||
"ollama_chat/deepseek-r1:70b": {"supports_function_calling": True},
|
||||
"ollama_chat/deepseek-r1:671b": {"supports_function_calling": True},
|
||||
"ollama_chat/deepseek-v3.1:latest": {"supports_function_calling": True},
|
||||
"ollama_chat/deepseek-v3.1:671b": {"supports_function_calling": True},
|
||||
"ollama_chat/deepseek-v3.1:671b-cloud": {"supports_function_calling": True},
|
||||
"ollama/deepseek-r1:latest": {"supports_function_calling": True},
|
||||
"ollama/deepseek-r1:1.5b": {"supports_function_calling": True},
|
||||
"ollama/deepseek-r1:7b": {"supports_function_calling": True},
|
||||
"ollama/deepseek-r1:8b": {"supports_function_calling": True},
|
||||
"ollama/deepseek-r1:14b": {"supports_function_calling": True},
|
||||
"ollama/deepseek-r1:32b": {"supports_function_calling": True},
|
||||
"ollama/deepseek-r1:70b": {"supports_function_calling": True},
|
||||
"ollama/deepseek-r1:671b": {"supports_function_calling": True},
|
||||
"ollama/deepseek-v3.1:latest": {"supports_function_calling": True},
|
||||
"ollama/deepseek-v3.1:671b": {"supports_function_calling": True},
|
||||
"ollama/deepseek-v3.1:671b-cloud": {"supports_function_calling": True},
|
||||
# Gemma3 models
|
||||
"ollama_chat/gemma3:latest": {"supports_function_calling": True},
|
||||
"ollama_chat/gemma3:270m": {"supports_function_calling": True},
|
||||
"ollama_chat/gemma3:1b": {"supports_function_calling": True},
|
||||
"ollama_chat/gemma3:4b": {"supports_function_calling": True},
|
||||
"ollama_chat/gemma3:12b": {"supports_function_calling": True},
|
||||
"ollama_chat/gemma3:27b": {"supports_function_calling": True},
|
||||
"ollama/gemma3:latest": {"supports_function_calling": True},
|
||||
"ollama/gemma3:270m": {"supports_function_calling": True},
|
||||
"ollama/gemma3:1b": {"supports_function_calling": True},
|
||||
"ollama/gemma3:4b": {"supports_function_calling": True},
|
||||
"ollama/gemma3:12b": {"supports_function_calling": True},
|
||||
"ollama/gemma3:27b": {"supports_function_calling": True},
|
||||
# Qwen models
|
||||
"ollama_chat/qwen3-coder:latest": {"supports_function_calling": True},
|
||||
"ollama_chat/qwen3-coder:30b": {"supports_function_calling": True},
|
||||
"ollama_chat/qwen3-coder:480b": {"supports_function_calling": True},
|
||||
"ollama_chat/qwen3-coder:480b-cloud": {"supports_function_calling": True},
|
||||
"ollama_chat/qwen3-vl:latest": {"supports_function_calling": True},
|
||||
"ollama_chat/qwen3-vl:2b": {"supports_function_calling": True},
|
||||
"ollama_chat/qwen3-vl:4b": {"supports_function_calling": True},
|
||||
"ollama_chat/qwen3-vl:8b": {"supports_function_calling": True},
|
||||
"ollama_chat/qwen3-vl:30b": {"supports_function_calling": True},
|
||||
"ollama_chat/qwen3-vl:32b": {"supports_function_calling": True},
|
||||
"ollama_chat/qwen3-vl:235b": {"supports_function_calling": True},
|
||||
"ollama_chat/qwen3-vl:235b-cloud": {"supports_function_calling": True},
|
||||
"ollama_chat/qwen3-vl:235b-instruct-cloud": {
|
||||
"supports_function_calling": True
|
||||
},
|
||||
"ollama/qwen3-coder:latest": {"supports_function_calling": True},
|
||||
"ollama/qwen3-coder:30b": {"supports_function_calling": True},
|
||||
"ollama/qwen3-coder:480b": {"supports_function_calling": True},
|
||||
"ollama/qwen3-coder:480b-cloud": {"supports_function_calling": True},
|
||||
"ollama/qwen3-vl:latest": {"supports_function_calling": True},
|
||||
"ollama/qwen3-vl:2b": {"supports_function_calling": True},
|
||||
"ollama/qwen3-vl:4b": {"supports_function_calling": True},
|
||||
"ollama/qwen3-vl:8b": {"supports_function_calling": True},
|
||||
"ollama/qwen3-vl:30b": {"supports_function_calling": True},
|
||||
"ollama/qwen3-vl:32b": {"supports_function_calling": True},
|
||||
"ollama/qwen3-vl:235b": {"supports_function_calling": True},
|
||||
"ollama/qwen3-vl:235b-cloud": {"supports_function_calling": True},
|
||||
"ollama/qwen3-vl:235b-instruct-cloud": {"supports_function_calling": True},
|
||||
# Kimi
|
||||
"ollama_chat/kimi-k2:1t": {"supports_function_calling": True},
|
||||
"ollama_chat/kimi-k2:1t-cloud": {"supports_function_calling": True},
|
||||
"ollama/kimi-k2:1t": {"supports_function_calling": True},
|
||||
"ollama/kimi-k2:1t-cloud": {"supports_function_calling": True},
|
||||
# GLM
|
||||
"ollama_chat/glm-4.6:cloud": {"supports_function_calling": True},
|
||||
"ollama_chat/glm-4.6": {"supports_function_calling": True},
|
||||
"ollama/glm-4.6": {"supports_function_calling": True},
|
||||
"ollama/glm-4.6-cloud": {"supports_function_calling": True},
|
||||
}
|
||||
)
|
||||
model_cost: dict[str, dict[str, bool]] = {
|
||||
# GPT-OSS models
|
||||
"ollama_chat/gpt-oss:120b-cloud": {"supports_function_calling": True},
|
||||
"ollama_chat/gpt-oss:120b": {"supports_function_calling": True},
|
||||
"ollama_chat/gpt-oss:20b-cloud": {"supports_function_calling": True},
|
||||
"ollama_chat/gpt-oss:20b": {"supports_function_calling": True},
|
||||
"ollama/gpt-oss:120b-cloud": {"supports_function_calling": True},
|
||||
"ollama/gpt-oss:120b": {"supports_function_calling": True},
|
||||
"ollama/gpt-oss:20b-cloud": {"supports_function_calling": True},
|
||||
"ollama/gpt-oss:20b": {"supports_function_calling": True},
|
||||
# DeepSeek models
|
||||
"ollama_chat/deepseek-r1:latest": {"supports_function_calling": True},
|
||||
"ollama_chat/deepseek-r1:1.5b": {"supports_function_calling": True},
|
||||
"ollama_chat/deepseek-r1:7b": {"supports_function_calling": True},
|
||||
"ollama_chat/deepseek-r1:8b": {"supports_function_calling": True},
|
||||
"ollama_chat/deepseek-r1:14b": {"supports_function_calling": True},
|
||||
"ollama_chat/deepseek-r1:32b": {"supports_function_calling": True},
|
||||
"ollama_chat/deepseek-r1:70b": {"supports_function_calling": True},
|
||||
"ollama_chat/deepseek-r1:671b": {"supports_function_calling": True},
|
||||
"ollama_chat/deepseek-v3.1:latest": {"supports_function_calling": True},
|
||||
"ollama_chat/deepseek-v3.1:671b": {"supports_function_calling": True},
|
||||
"ollama_chat/deepseek-v3.1:671b-cloud": {"supports_function_calling": True},
|
||||
"ollama/deepseek-r1:latest": {"supports_function_calling": True},
|
||||
"ollama/deepseek-r1:1.5b": {"supports_function_calling": True},
|
||||
"ollama/deepseek-r1:7b": {"supports_function_calling": True},
|
||||
"ollama/deepseek-r1:8b": {"supports_function_calling": True},
|
||||
"ollama/deepseek-r1:14b": {"supports_function_calling": True},
|
||||
"ollama/deepseek-r1:32b": {"supports_function_calling": True},
|
||||
"ollama/deepseek-r1:70b": {"supports_function_calling": True},
|
||||
"ollama/deepseek-r1:671b": {"supports_function_calling": True},
|
||||
"ollama/deepseek-v3.1:latest": {"supports_function_calling": True},
|
||||
"ollama/deepseek-v3.1:671b": {"supports_function_calling": True},
|
||||
"ollama/deepseek-v3.1:671b-cloud": {"supports_function_calling": True},
|
||||
# Gemma3 models
|
||||
"ollama_chat/gemma3:latest": {"supports_function_calling": True},
|
||||
"ollama_chat/gemma3:270m": {"supports_function_calling": True},
|
||||
"ollama_chat/gemma3:1b": {"supports_function_calling": True},
|
||||
"ollama_chat/gemma3:4b": {"supports_function_calling": True},
|
||||
"ollama_chat/gemma3:12b": {"supports_function_calling": True},
|
||||
"ollama_chat/gemma3:27b": {"supports_function_calling": True},
|
||||
"ollama/gemma3:latest": {"supports_function_calling": True},
|
||||
"ollama/gemma3:270m": {"supports_function_calling": True},
|
||||
"ollama/gemma3:1b": {"supports_function_calling": True},
|
||||
"ollama/gemma3:4b": {"supports_function_calling": True},
|
||||
"ollama/gemma3:12b": {"supports_function_calling": True},
|
||||
"ollama/gemma3:27b": {"supports_function_calling": True},
|
||||
# Qwen models
|
||||
"ollama_chat/qwen3-coder:latest": {"supports_function_calling": True},
|
||||
"ollama_chat/qwen3-coder:30b": {"supports_function_calling": True},
|
||||
"ollama_chat/qwen3-coder:480b": {"supports_function_calling": True},
|
||||
"ollama_chat/qwen3-coder:480b-cloud": {"supports_function_calling": True},
|
||||
"ollama_chat/qwen3-vl:latest": {"supports_function_calling": True},
|
||||
"ollama_chat/qwen3-vl:2b": {"supports_function_calling": True},
|
||||
"ollama_chat/qwen3-vl:4b": {"supports_function_calling": True},
|
||||
"ollama_chat/qwen3-vl:8b": {"supports_function_calling": True},
|
||||
"ollama_chat/qwen3-vl:30b": {"supports_function_calling": True},
|
||||
"ollama_chat/qwen3-vl:32b": {"supports_function_calling": True},
|
||||
"ollama_chat/qwen3-vl:235b": {"supports_function_calling": True},
|
||||
"ollama_chat/qwen3-vl:235b-cloud": {"supports_function_calling": True},
|
||||
"ollama_chat/qwen3-vl:235b-instruct-cloud": {"supports_function_calling": True},
|
||||
"ollama/qwen3-coder:latest": {"supports_function_calling": True},
|
||||
"ollama/qwen3-coder:30b": {"supports_function_calling": True},
|
||||
"ollama/qwen3-coder:480b": {"supports_function_calling": True},
|
||||
"ollama/qwen3-coder:480b-cloud": {"supports_function_calling": True},
|
||||
"ollama/qwen3-vl:latest": {"supports_function_calling": True},
|
||||
"ollama/qwen3-vl:2b": {"supports_function_calling": True},
|
||||
"ollama/qwen3-vl:4b": {"supports_function_calling": True},
|
||||
"ollama/qwen3-vl:8b": {"supports_function_calling": True},
|
||||
"ollama/qwen3-vl:30b": {"supports_function_calling": True},
|
||||
"ollama/qwen3-vl:32b": {"supports_function_calling": True},
|
||||
"ollama/qwen3-vl:235b": {"supports_function_calling": True},
|
||||
"ollama/qwen3-vl:235b-cloud": {"supports_function_calling": True},
|
||||
"ollama/qwen3-vl:235b-instruct-cloud": {"supports_function_calling": True},
|
||||
# Kimi
|
||||
"ollama_chat/kimi-k2:1t": {"supports_function_calling": True},
|
||||
"ollama_chat/kimi-k2:1t-cloud": {"supports_function_calling": True},
|
||||
"ollama/kimi-k2:1t": {"supports_function_calling": True},
|
||||
"ollama/kimi-k2:1t-cloud": {"supports_function_calling": True},
|
||||
# GLM
|
||||
"ollama_chat/glm-4.6:cloud": {"supports_function_calling": True},
|
||||
"ollama_chat/glm-4.6": {"supports_function_calling": True},
|
||||
"ollama/glm-4.6": {"supports_function_calling": True},
|
||||
"ollama/glm-4.6-cloud": {"supports_function_calling": True},
|
||||
}
|
||||
|
||||
# Seed the keys first. litellm treats a model it does not already know as a
|
||||
# live Ollama deployment: it POSTs to the Ollama server for each name, then
|
||||
# stores the result under the prefix-stripped key. One pass would leave
|
||||
# every ollama_chat/* model without function calling.
|
||||
for model_name in model_cost:
|
||||
litellm.model_cost.setdefault(model_name, {})
|
||||
litellm.register_model(model_cost=model_cost)
|
||||
|
||||
|
||||
def load_model_metadata_enrichments() -> None:
|
||||
@@ -164,7 +169,25 @@ def load_model_metadata_enrichments() -> None:
|
||||
logger.error("Failed to load model metadata enrichments: %s", e)
|
||||
|
||||
|
||||
_INIT_LOCK = threading.Lock()
|
||||
_initialized = False
|
||||
|
||||
|
||||
def initialize_litellm() -> None:
|
||||
configure_litellm_settings()
|
||||
register_ollama_models()
|
||||
load_model_metadata_enrichments()
|
||||
"""Configure the process-wide litellm module. Safe to call from any thread.
|
||||
|
||||
Runs once per process. Repeat calls are not free: they re-register every
|
||||
Ollama model, which rebuilds litellm's model lookup maps and clears its
|
||||
model-info caches, and they mutate litellm.model_cost while other threads
|
||||
iterate it.
|
||||
"""
|
||||
global _initialized
|
||||
if _initialized:
|
||||
return
|
||||
with _INIT_LOCK:
|
||||
if _initialized:
|
||||
return
|
||||
configure_litellm_settings()
|
||||
register_ollama_models()
|
||||
load_model_metadata_enrichments()
|
||||
_initialized = True
|
||||
@@ -0,0 +1,133 @@
|
||||
import copy
|
||||
import threading
|
||||
from types import ModuleType
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
_SAMPLE_MODELS = (
|
||||
"ollama_chat/gpt-oss:20b",
|
||||
"ollama_chat/deepseek-r1:14b",
|
||||
"ollama/glm-4.6",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def configured_litellm() -> ModuleType:
|
||||
"""The litellm module after the singleton package has configured it."""
|
||||
import litellm
|
||||
|
||||
import onyx.llm.litellm_singleton # noqa: F401
|
||||
|
||||
return litellm
|
||||
|
||||
|
||||
def test_ollama_chat_models_support_function_calling(
|
||||
configured_litellm: ModuleType,
|
||||
) -> None:
|
||||
for model in (
|
||||
"ollama_chat/gpt-oss:20b",
|
||||
"ollama_chat/deepseek-r1:14b",
|
||||
"ollama_chat/glm-4.6",
|
||||
"ollama/gpt-oss:20b",
|
||||
):
|
||||
assert configured_litellm.supports_function_calling(model=model) is True
|
||||
|
||||
|
||||
def test_register_ollama_models_needs_one_pass_and_no_server(
|
||||
configured_litellm: ModuleType,
|
||||
) -> None:
|
||||
"""Registration must not depend on a reachable Ollama server, or on a second pass.
|
||||
|
||||
litellm treats an unknown model name as a live deployment and probes the
|
||||
Ollama server for it, then stores the entry under the prefix-stripped name.
|
||||
Seeding the keys keeps it on the static path.
|
||||
"""
|
||||
from onyx.llm.litellm_singleton.config import register_ollama_models
|
||||
|
||||
# The names litellm checks before deciding a model is a live deployment.
|
||||
keys = {
|
||||
name
|
||||
for model in _SAMPLE_MODELS
|
||||
for name in (
|
||||
model,
|
||||
model.split("/", 1)[1],
|
||||
"ollama/" + model.split("/", 1)[1],
|
||||
"ollama_chat/" + model.split("/", 1)[1],
|
||||
)
|
||||
}
|
||||
snapshot = {
|
||||
key: copy.deepcopy(configured_litellm.model_cost[key])
|
||||
for key in keys
|
||||
if key in configured_litellm.model_cost
|
||||
}
|
||||
|
||||
try:
|
||||
for key in keys:
|
||||
configured_litellm.model_cost.pop(key, None)
|
||||
|
||||
with patch.object(
|
||||
configured_litellm.module_level_client,
|
||||
"post",
|
||||
side_effect=AssertionError("registration probed the Ollama server"),
|
||||
):
|
||||
register_ollama_models()
|
||||
|
||||
for model in _SAMPLE_MODELS:
|
||||
assert configured_litellm.supports_function_calling(model=model) is True
|
||||
finally:
|
||||
configured_litellm.model_cost.update(snapshot)
|
||||
register_ollama_models()
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("configured_litellm")
|
||||
def test_initialize_litellm_runs_once() -> None:
|
||||
from onyx.llm.litellm_singleton import config
|
||||
|
||||
with (
|
||||
patch.object(config, "configure_litellm_settings") as configure,
|
||||
patch.object(config, "register_ollama_models") as register,
|
||||
patch.object(config, "load_model_metadata_enrichments") as enrich,
|
||||
):
|
||||
# Importing the package already initialized litellm. Reset so this
|
||||
# covers the first call and the second one, not just the second.
|
||||
config._initialized = False
|
||||
try:
|
||||
config.initialize_litellm()
|
||||
config.initialize_litellm()
|
||||
finally:
|
||||
config._initialized = True
|
||||
|
||||
assert configure.call_count == 1
|
||||
assert register.call_count == 1
|
||||
assert enrich.call_count == 1
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("configured_litellm")
|
||||
def test_initialize_litellm_runs_once_under_concurrency() -> None:
|
||||
from onyx.llm.litellm_singleton import config
|
||||
|
||||
with (
|
||||
patch.object(config, "configure_litellm_settings") as configure,
|
||||
patch.object(config, "register_ollama_models") as register,
|
||||
patch.object(config, "load_model_metadata_enrichments") as enrich,
|
||||
):
|
||||
config._initialized = False
|
||||
try:
|
||||
barrier = threading.Barrier(8)
|
||||
|
||||
def worker() -> None:
|
||||
barrier.wait()
|
||||
config.initialize_litellm()
|
||||
|
||||
threads = [threading.Thread(target=worker) for _ in range(8)]
|
||||
for thread in threads:
|
||||
thread.start()
|
||||
for thread in threads:
|
||||
thread.join()
|
||||
finally:
|
||||
config._initialized = True
|
||||
|
||||
assert configure.call_count == 1
|
||||
assert register.call_count == 1
|
||||
assert enrich.call_count == 1
|
||||
Reference in new issue
Block a user