-
Notifications
You must be signed in to change notification settings - Fork 3.4k
prewarm LLM to skip connection establishment #6484
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -102,6 +102,7 @@ class LLM( | |
| def __init__(self) -> None: | ||
| super().__init__() | ||
| self._label = f"{type(self).__module__}.{type(self).__name__}" | ||
| self._prewarm_task: asyncio.Task[None] | None = None | ||
|
|
||
| @property | ||
| def label(self) -> str: | ||
|
|
@@ -144,10 +145,47 @@ def chat( | |
| ) -> LLMStream: ... | ||
|
|
||
| def prewarm(self) -> None: | ||
| """Pre-warm connection to the LLM service""" | ||
| pass | ||
| """Pre-warm connection to the LLM service. | ||
|
|
||
| async def aclose(self) -> None: ... | ||
| Establishes DNS resolution and the TLS connection to the provider before the | ||
| first inference request, reducing time-to-first-token on the initial reply. | ||
| It is called automatically when an ``AgentSession`` is constructed and when an | ||
| agent activity starts, but can also be called directly. | ||
|
|
||
| Non-blocking (fire-and-forget), idempotent, and a no-op when no event loop | ||
| is running. Providers enable it by overriding ``_prewarm_impl``. | ||
| """ | ||
| if type(self)._prewarm_impl is LLM._prewarm_impl: | ||
| return # no provider-specific prewarm implemented | ||
|
|
||
| if self._prewarm_task is not None: | ||
| return | ||
|
|
||
| try: | ||
| asyncio.get_running_loop() | ||
| except RuntimeError: | ||
| return # no running event loop, prewarm will be retried when the agent starts | ||
|
|
||
| async def _prewarm() -> None: | ||
| try: | ||
| await self._prewarm_impl() | ||
| except Exception: | ||
| pass | ||
|
|
||
| self._prewarm_task = asyncio.create_task(_prewarm()) | ||
|
|
||
| async def _prewarm_impl(self) -> None: | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Q: should we add support for FallbackAdapter? |
||
| """Provider-specific prewarm request, overriding it enables ``prewarm()``. | ||
|
|
||
| Implementations should perform a cheap, token-free request (e.g. listing | ||
| models) using the same client that serves chat requests, so DNS + TLS are | ||
| established and a keep-alive connection is left in the pool. Exceptions | ||
| are swallowed by ``prewarm()``; subclasses overriding ``aclose`` must call | ||
| ``await super().aclose()`` to cancel an in-flight prewarm.""" | ||
|
|
||
| async def aclose(self) -> None: | ||
| if self._prewarm_task is not None: | ||
| await aio.cancel_and_wait(self._prewarm_task) | ||
|
|
||
| async def __aenter__(self) -> LLM: | ||
| return self | ||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -30,7 +30,7 @@ | |
| from .. import cli, inference, llm, stt, tts, utils, vad | ||
| from .._exceptions import APIError | ||
| from ..job import get_job_context | ||
| from ..llm import AgentHandoff, ChatContext, MetricsReport | ||
| from ..llm import LLM, AgentHandoff, ChatContext, MetricsReport | ||
| from ..llm.chat_context import Instructions | ||
| from ..log import logger | ||
| from ..metrics import AgentSessionUsage, ModelUsageCollector | ||
|
|
@@ -477,6 +477,15 @@ def __init__( | |
| self._llm = llm or None | ||
| self._tts = tts or None | ||
|
|
||
| # eagerly establish DNS/TLS to the LLM provider so the first inference | ||
| # request doesn't pay connection setup costs | ||
| if isinstance(self._llm, LLM): | ||
| try: | ||
| asyncio.get_running_loop() | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. There is always an event loop
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Adding to this, I think we can skip checking it here since _llm.prewarm already does that internally. |
||
| self._llm.prewarm() | ||
| except RuntimeError: | ||
| pass # no running event loop, prewarm runs at agent activity start instead | ||
|
|
||
| self._keyterm_detector = KeytermDetector( | ||
| static_keyterms=self._opts.stt_context_options["keyterms"], | ||
| options=self._opts.stt_context_options["keyterm_detection"], | ||
|
|
||
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Should we clear the
_prewarm_taskhere?