diff --git a/private_gpt/components/tools/tool_scheduler.py b/private_gpt/components/tools/tool_scheduler.py index 0f17d96e..3ab811de 100644 --- a/private_gpt/components/tools/tool_scheduler.py +++ b/private_gpt/components/tools/tool_scheduler.py @@ -1,14 +1,17 @@ from __future__ import annotations import asyncio +import dataclasses import logging +import time import uuid from abc import ABC, abstractmethod -from collections.abc import Awaitable, Callable -from typing import Any +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from collections.abc import Awaitable, Callable from injector import inject, singleton -from pydantic import BaseModel, ConfigDict from private_gpt.components.tools.tool_names import ( BASH_TOOL_NAME, @@ -64,9 +67,8 @@ class ImmediateToolScheduler(BaseToolScheduler): return await func() -class _PendingCall(BaseModel): - model_config = ConfigDict(arbitrary_types_allowed=True) - +@dataclasses.dataclass +class _PendingCall: score: float counter: int entry_id: str @@ -187,3 +189,69 @@ class QueuedToolScheduler(BaseToolScheduler): if not future.done(): future.cancel() raise + + +@singleton +class AdaptiveToolScheduler(BaseToolScheduler): + @inject + def __init__( + self, + immediate: ImmediateToolScheduler, + queued: QueuedToolScheduler, + settings: Settings, + ) -> None: + cfg = settings.tool_scheduler.adaptive + self._immediate = immediate + self._queued = queued + self._threshold_ms: float = cfg.threshold_ms + self._recovery_ms: float = cfg.threshold_ms * cfg.recovery_ratio + self._alpha: float = cfg.ewma_alpha + self._ewma_ms: float | None = None + self._use_queue: bool = False + + def _update(self, elapsed_ms: float) -> None: + self._ewma_ms = ( + elapsed_ms + if self._ewma_ms is None + else self._alpha * elapsed_ms + (1 - self._alpha) * self._ewma_ms + ) + if not self._use_queue and self._ewma_ms > self._threshold_ms: + self._use_queue = True + logger.info("AdaptiveToolScheduler: -> queued (ewma=%.0fms)", self._ewma_ms) + elif self._use_queue and self._ewma_ms < self._recovery_ms: + self._use_queue = False + logger.info( + "AdaptiveToolScheduler: -> immediate (ewma=%.0fms)", self._ewma_ms + ) + + async def execute( + self, + tool_name: str, + chat_priority: int | None, + func: Callable[[], Awaitable[Any]], + ) -> Any: + scheduler = self._queued if self._use_queue else self._immediate + start = time.monotonic() + try: + return await scheduler.execute(tool_name, chat_priority, func) + finally: + self._update((time.monotonic() - start) * 1000) + + +@singleton +class ToolSchedulerFactory: + @inject + def __init__( + self, + settings: Settings, + queued: QueuedToolScheduler, + adaptive: AdaptiveToolScheduler, + ) -> None: + self._scheduler: BaseToolScheduler = { + "immediate": ImmediateToolScheduler(), + "queued": queued, + "adaptive": adaptive, + }[settings.tool_scheduler.mode] + + def get(self) -> BaseToolScheduler: + return self._scheduler diff --git a/private_gpt/components/vector_store/qdrant_client_builder.py b/private_gpt/components/vector_store/qdrant_client_builder.py index eab5d304..4635ed5f 100644 --- a/private_gpt/components/vector_store/qdrant_client_builder.py +++ b/private_gpt/components/vector_store/qdrant_client_builder.py @@ -11,7 +11,6 @@ from qdrant_client import ( # type: ignore[import-not-found] from private_gpt.settings.settings import Settings logger = logging.getLogger(__name__) -logger.setLevel(logging.DEBUG) class QdrantClients: diff --git a/private_gpt/server/chat/chat_service.py b/private_gpt/server/chat/chat_service.py index 9e038f4c..dfe6c91c 100644 --- a/private_gpt/server/chat/chat_service.py +++ b/private_gpt/server/chat/chat_service.py @@ -23,11 +23,7 @@ from private_gpt.components.llm.custom.base import ZylonLLM from private_gpt.components.llm.llm_component import LLMComponent from private_gpt.components.llm.models import ReasoningEffort from private_gpt.components.node_store.node_store_component import NodeStoreComponent -from private_gpt.components.tools.tool_scheduler import ( - BaseToolScheduler, - ImmediateToolScheduler, - QueuedToolScheduler, -) +from private_gpt.components.tools.tool_scheduler import ToolSchedulerFactory from private_gpt.components.vector_store.vector_store_component import ( VectorStoreComponent, ) @@ -155,7 +151,7 @@ class ChatService: chat_interceptor_service: ChatInterceptorService, models_service: ModelsService, container_registry: ContainerRegistry, - tool_scheduler: QueuedToolScheduler, + scheduler_factory: ToolSchedulerFactory, ) -> None: self.settings = settings self.llm_component = llm_component @@ -166,11 +162,7 @@ class ChatService: self.chat_interceptor_service = chat_interceptor_service self.models_service = models_service self.container_registry = container_registry - self._tool_scheduler: BaseToolScheduler = ( - tool_scheduler - if settings.tool_scheduler.enabled - else ImmediateToolScheduler() - ) + self._tool_scheduler = scheduler_factory.get() def _build_loop_engine(self) -> ChatLoopEngine: # Don't build a singleton since the interceptors diff --git a/private_gpt/settings/settings.py b/private_gpt/settings/settings.py index 4fccdf4f..6b0df8bc 100644 --- a/private_gpt/settings/settings.py +++ b/private_gpt/settings/settings.py @@ -1639,41 +1639,56 @@ class ToolSchedulerWeightsSettings(BaseModel): default=0.4, description="Weight given to the chat request priority signal (0-1).", ) - age: float = Field( - default=0.4, - description=( - "Weight given to wait time in the queue. " - "Higher values reduce starvation of long-waiting calls." - ), - ) complexity: float = Field( default=0.2, description="Weight given to the estimated tool complexity (0-1).", ) -class ToolSchedulerSettings(BaseModel): - enabled: bool = Field( - default=False, +class AdaptiveToolSchedulerSettings(BaseModel): + threshold_ms: float = Field( + default=2000.0, description=( - "When True, all tool executions are routed through a shared priority queue " - "that limits global concurrency and orders calls by urgency." + "EWMA of tool execution time (ms) above which the adaptive scheduler " + "switches from immediate to queued mode." + ), + ) + recovery_ratio: float = Field( + default=0.7, + description=( + "Fraction of threshold_ms at which the adaptive scheduler switches back " + "to immediate mode. Must be < 1 to provide hysteresis." + ), + ) + ewma_alpha: float = Field( + default=0.3, + description=( + "Smoothing factor for the exponentially weighted moving average. " + "Higher values react faster to load changes." + ), + ) + + +class ToolSchedulerSettings(BaseModel): + mode: Literal["immediate", "queued", "adaptive"] = Field( + default="immediate", + description=( + "'immediate' runs tools without queuing; " + "'queued' routes all calls through a shared priority queue; " + "'adaptive' switches between the two based on measured execution time." ), ) max_concurrent_tools: int = Field( default=5, - description="Maximum number of tool calls executed simultaneously across all chats.", - ) - max_age_seconds: float = Field( - default=60.0, - description=( - "Time (seconds) after which a queued tool call reaches maximum age urgency. " - "Prevents starvation by boosting priority of long-waiting calls." - ), + description="Maximum simultaneous tool calls (used by 'queued' and 'adaptive' modes).", ) weights: ToolSchedulerWeightsSettings = Field( default_factory=ToolSchedulerWeightsSettings, - description="Relative weights for the three priority signals.", + description="Priority weights for the queued scheduler.", + ) + adaptive: AdaptiveToolSchedulerSettings = Field( + default_factory=AdaptiveToolSchedulerSettings, + description="Settings for the adaptive scheduler.", ) diff --git a/settings.yaml b/settings.yaml index fc6be972..4afa1e91 100644 --- a/settings.yaml +++ b/settings.yaml @@ -290,13 +290,15 @@ semaphore: mode: ${PGPT_SEMAPHORE_MODE:memory} tool_scheduler: - enabled: ${PGPT_TOOL_SCHEDULER_ENABLED:false} - max_concurrent_tools: ${PGPT_TOOL_SCHEDULER_MAX_CONCURRENT_TOOLS:5} - max_age_seconds: ${PGPT_TOOL_SCHEDULER_MAX_AGE_SECONDS:60} + mode: ${PGPT_TOOL_SCHEDULER_MODE:adaptive} + max_concurrent_tools: ${PGPT_TOOL_SCHEDULER_MAX_CONCURRENT_TOOLS:4} weights: chat_priority: ${PGPT_TOOL_SCHEDULER_WEIGHT_PRIORITY:0.4} - age: ${PGPT_TOOL_SCHEDULER_WEIGHT_AGE:0.4} complexity: ${PGPT_TOOL_SCHEDULER_WEIGHT_COMPLEXITY:0.2} + adaptive: + threshold_ms: ${PGPT_TOOL_SCHEDULER_ADAPTIVE_THRESHOLD_MS:800} + recovery_ratio: ${PGPT_TOOL_SCHEDULER_ADAPTIVE_RECOVERY_RATIO:0.7} + ewma_alpha: ${PGPT_TOOL_SCHEDULER_ADAPTIVE_EWMA_ALPHA:0.5} transformation: pptx: