Use async for model metrics initialization#294
Use async for model metrics initialization#294omertuc wants to merge 1 commit intolightspeed-core:mainfrom
Conversation
`LlamaStackClientHolder` is sync, but internally within llama-stack it
glues into the async client, by instantiating its own event loop. This
somehow interferes with the async event loop of uvicorn:
<details>
<summary>Traceback</summary>
```
Traceback (most recent call last):
File "/app-root/src/lightspeed_stack.py", line 91, in <module>
main()
File "/app-root/src/lightspeed_stack.py", line 86, in main
start_uvicorn(configuration.service_configuration)
File "/app-root/src/runners/uvicorn.py", line 20, in start_uvicorn
uvicorn.run(
File "/app-root/.venv/lib64/python3.12/site-packages/uvicorn/main.py", line 580, in run
server.run()
File "/app-root/.venv/lib64/python3.12/site-packages/uvicorn/server.py", line 67, in run
return asyncio.run(self.serve(sockets=sockets))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/lib64/python3.12/asyncio/runners.py", line 195, in run
return runner.run(main)
^^^^^^^^^^^^^^^^
File "/usr/lib64/python3.12/asyncio/runners.py", line 118, in run
return self._loop.run_until_complete(task)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/lib64/python3.12/asyncio/base_events.py", line 691, in run_until_complete
return future.result()
^^^^^^^^^^^^^^^
File "/app-root/.venv/lib64/python3.12/site-packages/uvicorn/server.py", line 71, in serve
await self._serve(sockets)
File "/app-root/.venv/lib64/python3.12/site-packages/uvicorn/server.py", line 78, in _serve
config.load()
File "/app-root/.venv/lib64/python3.12/site-packages/uvicorn/config.py", line 436, in load
self.loaded_app = import_from_string(self.app)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/app-root/.venv/lib64/python3.12/site-packages/uvicorn/importer.py", line 19, in import_from_string
module = importlib.import_module(module_str)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/lib64/python3.12/importlib/__init__.py", line 90, in import_module
return _bootstrap._gcd_import(name[level:], package, level)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "<frozen importlib._bootstrap>", line 1387, in _gcd_import
File "<frozen importlib._bootstrap>", line 1360, in _find_and_load
File "<frozen importlib._bootstrap>", line 1331, in _find_and_load_unlocked
File "<frozen importlib._bootstrap>", line 935, in _load_unlocked
File "<frozen importlib._bootstrap_external>", line 999, in exec_module
File "<frozen importlib._bootstrap>", line 488, in _call_with_frames_removed
File "/app-root/src/app/main.py", line 77, in <module>
setup_model_metrics()
File "/app-root/src/metrics/utils.py", line 16, in setup_model_metrics
for model in client.models.list()
^^^^^^^^^^^^^^^^^^^^
File "/app-root/.venv/lib64/python3.12/site-packages/llama_stack_client/resources/models.py", line 93, in list
return self._get(
^^^^^^^^^^
File "/app-root/.venv/lib64/python3.12/site-packages/llama_stack_client/_base_client.py", line 1178, in get
return cast(ResponseT, self.request(cast_to, opts, stream=stream, stream_cls=stream_cls))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/app-root/llama-stack/llama_stack/distribution/library_client.py", line 177, in request
result = loop.run_until_complete(self.async_client.request(*args, **kwargs))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/lib64/python3.12/asyncio/base_events.py", line 667, in run_until_complete
self._check_running()
File "/usr/lib64/python3.12/asyncio/base_events.py", line 628, in _check_running
raise RuntimeError(
RuntimeError: Cannot run the event loop while another loop is running
sys:1: RuntimeWarning: coroutine 'AsyncLlamaStackAsLibraryClient.request' was never awaited
```
</details>
Not sure why this was not caught earlier, maybe it only happens when
running llama-stack in library mode?
I've changed the `setup_model_metrics` function to be async and use
`AsyncLlamaStackClientHolder` instead of `LlamaStackClientHolder`. It
also means we should call it from `startup_event`
WalkthroughThe changes refactor the model metrics setup process to be asynchronous. The Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~8 minutes Poem
✨ Finishing Touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
src/app/main.py (1)
81-86: Consider error handling for startup operationsWhile the async integration is correct, consider adding error handling around
setup_model_metrics()to prevent startup failures from completely breaking the application. The current implementation may cause the entire application to fail to start if model metrics setup encounters issues.Consider wrapping the call in a try-except block:
async def startup_event() -> None: """Perform logger setup and other initializations on service startup.""" try: await setup_model_metrics() except Exception as e: logger.error("Failed to setup model metrics: %s", e) # Consider whether to continue startup or fail fast based on requirements logger.info("Registering MCP servers") await register_mcp_servers_async(logger, configuration.configuration) get_logger("app.endpoints.handlers") logger.info("App startup complete")
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
src/app/main.py(1 hunks)src/metrics/utils.py(1 hunks)
🔇 Additional comments (5)
src/metrics/utils.py (3)
4-4: LGTM: Correct import for async patternThe import change from
LlamaStackClientHoldertoAsyncLlamaStackClientHolderaligns with the async conversion requirements.
11-11: LGTM: Function correctly converted to asyncThe function signature is properly updated to async, which is necessary to support the awaited client operations.
16-16: LGTM: Correct async pattern for model listingThe
await client.models.list()call is properly implemented for the async client pattern.src/app/main.py (2)
80-80: LGTM: Docstring accurately updatedThe docstring correctly reflects that the function now performs additional initializations beyond just logger setup.
81-81: LGTM: Proper async integration with startup eventThe
await setup_model_metrics()call is correctly placed in the async startup event, resolving the event loop conflict described in the PR objectives.
| async def setup_model_metrics() -> None: | ||
| """Perform setup of all metrics related to LLM model and provider.""" | ||
| client = LlamaStackClientHolder().get_client() | ||
| client = AsyncLlamaStackClientHolder().get_client() |
There was a problem hiding this comment.
💡 Verification agent
🧩 Analysis chain
Verify error handling for async client initialization
The async client initialization looks correct, but ensure that any potential exceptions from AsyncLlamaStackClientHolder().get_client() are properly handled upstream.
🏁 Script executed:
#!/bin/bash
# Description: Check if AsyncLlamaStackClientHolder.get_client() can raise exceptions that need handling
# Expected: Find the implementation and any exception handling patterns
ast-grep --pattern $'class AsyncLlamaStackClientHolder {
$$$
get_client($$$) {
$$$
}
$$$
}'Length of output: 104
🏁 Script executed:
#!/usr/bin/env bash
# Locate the AsyncLlamaStackClientHolder class and its get_client method
echo "==> Searching for class AsyncLlamaStackClientHolder definition"
rg -n "class AsyncLlamaStackClientHolder" .
echo
echo "==> Searching for get_client method definition"
rg -n "def get_client" .Length of output: 528
🏁 Script executed:
#!/usr/bin/env bash
# Display the AsyncLlamaStackClientHolder class and its get_client implementation
echo "==> Showing lines 45–120 of src/client.py for context"
sed -n '45,120p' src/client.pyLength of output: 1974
🏁 Script executed:
#!/usr/bin/env bash
# Display the first 200 lines of src/metrics/utils.py to review get_client usage
echo "==> Showing lines 1–200 of src/metrics/utils.py"
sed -n '1,200p' src/metrics/utils.pyLength of output: 1733
Handle potential RuntimeError from uninitialized async client
AsyncLlamaStackClientHolder.get_client() will raise a RuntimeError if load() hasn’t been called. In src/metrics/utils.py at line 13, wrap the call in a try/except or otherwise ensure the client is initialized before fetching metrics:
Possible fix:
diff --git a/src/metrics/utils.py b/src/metrics/utils.py
@@ async def setup_model_metrics() -> None:
- client = AsyncLlamaStackClientHolder().get_client()
+ try:
+ client = AsyncLlamaStackClientHolder().get_client()
+ except RuntimeError as e:
+ logger.error("Metrics setup skipped: %s", e)
+ returnThis will prevent an uncaught exception during metrics initialization.
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| client = AsyncLlamaStackClientHolder().get_client() | |
| try: | |
| client = AsyncLlamaStackClientHolder().get_client() | |
| except RuntimeError as e: | |
| logger.error("Metrics setup skipped: %s", e) | |
| return |
🤖 Prompt for AI Agents
In src/metrics/utils.py at line 13, the call to
AsyncLlamaStackClientHolder().get_client() can raise a RuntimeError if the
client is not initialized. To fix this, wrap the get_client() call in a
try/except block to catch RuntimeError and handle it gracefully, or ensure that
load() is called on AsyncLlamaStackClientHolder before calling get_client().
This will prevent uncaught exceptions during metrics initialization.
|
Oh are you kidding me... #293 |
2 participants
Description
LlamaStackClientHolderis sync, but internally within llama-stack it glues into the async client, by instantiating its own event loop. This somehow interferes with the async event loop of uvicorn:Traceback
Not sure why this was not caught earlier, maybe it only happens when running llama-stack in library mode?
I've changed the
setup_model_metricsfunction to be async and useAsyncLlamaStackClientHolderinstead ofLlamaStackClientHolder. It also means we should call it fromstartup_eventType of change
Related Tickets & Documents
Checklist before requesting a review
Testing
I've ran it manually in the same setup where the error occurs and made sure the error no longer occurs