forked from phoenix-oss/llama-stack-mirror
remove default logger handlers when using libcli with notebook (#718)
# What does this PR do? Remove the default log handlers for notebook to avoid polluting logs
This commit is contained in:
parent
e86271aeac
commit
0bc5d05243
1 changed files with 15 additions and 3 deletions
|
@ -7,6 +7,7 @@
|
||||||
import asyncio
|
import asyncio
|
||||||
import inspect
|
import inspect
|
||||||
import json
|
import json
|
||||||
|
import logging
|
||||||
import os
|
import os
|
||||||
import queue
|
import queue
|
||||||
import threading
|
import threading
|
||||||
|
@ -16,7 +17,6 @@ from pathlib import Path
|
||||||
from typing import Any, Generator, get_args, get_origin, Optional, TypeVar
|
from typing import Any, Generator, get_args, get_origin, Optional, TypeVar
|
||||||
|
|
||||||
import httpx
|
import httpx
|
||||||
|
|
||||||
import yaml
|
import yaml
|
||||||
from llama_stack_client import (
|
from llama_stack_client import (
|
||||||
APIResponse,
|
APIResponse,
|
||||||
|
@ -28,7 +28,6 @@ from llama_stack_client import (
|
||||||
)
|
)
|
||||||
from pydantic import BaseModel, TypeAdapter
|
from pydantic import BaseModel, TypeAdapter
|
||||||
from rich.console import Console
|
from rich.console import Console
|
||||||
|
|
||||||
from termcolor import cprint
|
from termcolor import cprint
|
||||||
|
|
||||||
from llama_stack.distribution.build import print_pip_install_help
|
from llama_stack.distribution.build import print_pip_install_help
|
||||||
|
@ -42,7 +41,6 @@ from llama_stack.distribution.stack import (
|
||||||
redact_sensitive_fields,
|
redact_sensitive_fields,
|
||||||
replace_env_vars,
|
replace_env_vars,
|
||||||
)
|
)
|
||||||
|
|
||||||
from llama_stack.providers.utils.telemetry.tracing import (
|
from llama_stack.providers.utils.telemetry.tracing import (
|
||||||
end_trace,
|
end_trace,
|
||||||
setup_logger,
|
setup_logger,
|
||||||
|
@ -174,6 +172,7 @@ class LlamaStackAsLibraryClient(LlamaStackClient):
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
config_path_or_template_name: str,
|
config_path_or_template_name: str,
|
||||||
|
skip_logger_removal: bool = False,
|
||||||
custom_provider_registry: Optional[ProviderRegistry] = None,
|
custom_provider_registry: Optional[ProviderRegistry] = None,
|
||||||
):
|
):
|
||||||
super().__init__()
|
super().__init__()
|
||||||
|
@ -181,15 +180,28 @@ class LlamaStackAsLibraryClient(LlamaStackClient):
|
||||||
config_path_or_template_name, custom_provider_registry
|
config_path_or_template_name, custom_provider_registry
|
||||||
)
|
)
|
||||||
self.pool_executor = ThreadPoolExecutor(max_workers=4)
|
self.pool_executor = ThreadPoolExecutor(max_workers=4)
|
||||||
|
self.skip_logger_removal = skip_logger_removal
|
||||||
|
|
||||||
def initialize(self):
|
def initialize(self):
|
||||||
if in_notebook():
|
if in_notebook():
|
||||||
import nest_asyncio
|
import nest_asyncio
|
||||||
|
|
||||||
nest_asyncio.apply()
|
nest_asyncio.apply()
|
||||||
|
if not self.skip_logger_removal:
|
||||||
|
self._remove_root_logger_handlers()
|
||||||
|
|
||||||
return asyncio.run(self.async_client.initialize())
|
return asyncio.run(self.async_client.initialize())
|
||||||
|
|
||||||
|
def _remove_root_logger_handlers(self):
|
||||||
|
"""
|
||||||
|
Remove all handlers from the root logger. Needed to avoid polluting the console with logs.
|
||||||
|
"""
|
||||||
|
root_logger = logging.getLogger()
|
||||||
|
|
||||||
|
for handler in root_logger.handlers[:]:
|
||||||
|
root_logger.removeHandler(handler)
|
||||||
|
print(f"Removed handler {handler.__class__.__name__} from root logger")
|
||||||
|
|
||||||
def _get_path(
|
def _get_path(
|
||||||
self,
|
self,
|
||||||
cast_to: Any,
|
cast_to: Any,
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue