diff --git a/llama_stack/distribution/datatypes.py b/llama_stack/distribution/datatypes.py index f62996081..7e1d8c016 100644 --- a/llama_stack/distribution/datatypes.py +++ b/llama_stack/distribution/datatypes.py @@ -117,6 +117,14 @@ class Provider(BaseModel): config: Dict[str, Any] +class LoggingConfig(BaseModel): + category_levels: Dict[str, str] = Field( + default_factory=Dict, + description=""" + Dictionary of different logging configurations for different portions (ex: core, server) of llama stack""", + ) + + class ServerConfig(BaseModel): port: int = Field( default=8321, @@ -176,6 +184,8 @@ a default SQLite store will be used.""", benchmarks: List[BenchmarkInput] = Field(default_factory=list) tool_groups: List[ToolGroupInput] = Field(default_factory=list) + logging: Optional[LoggingConfig] = Field(default=None, description="Configuration for Llama Stack Logging") + server: ServerConfig = Field( default_factory=ServerConfig, description="Configuration for the HTTP(S) server", diff --git a/llama_stack/distribution/server/server.py b/llama_stack/distribution/server/server.py index 8f9500ae9..b37b3a007 100644 --- a/llama_stack/distribution/server/server.py +++ b/llama_stack/distribution/server/server.py @@ -25,7 +25,7 @@ from fastapi.responses import JSONResponse, StreamingResponse from pydantic import BaseModel, ValidationError from typing_extensions import Annotated -from llama_stack.distribution.datatypes import StackRunConfig +from llama_stack.distribution.datatypes import LoggingConfig, StackRunConfig from llama_stack.distribution.distribution import builtin_automatically_routed_apis from llama_stack.distribution.request_headers import ( PROVIDER_DATA_VAR, @@ -306,34 +306,42 @@ def main(): args = parser.parse_args() - if args.env: - for env_pair in args.env: - try: - key, value = validate_env_pair(env_pair) - logger.info(f"Setting CLI environment variable {key} => {value}") - os.environ[key] = value - except ValueError as e: - logger.error(f"Error: {str(e)}") - sys.exit(1) - + log_line = "" if args.yaml_config: # if the user provided a config file, use it, even if template was specified config_file = Path(args.yaml_config) if not config_file.exists(): raise ValueError(f"Config file {config_file} does not exist") - logger.info(f"Using config file: {config_file}") + log_line = f"Using config file: {config_file}" elif args.template: config_file = Path(REPO_ROOT) / "llama_stack" / "templates" / args.template / "run.yaml" if not config_file.exists(): raise ValueError(f"Template {args.template} does not exist") - logger.info(f"Using template {args.template} config file: {config_file}") + log_line = f"Using template {args.template} config file: {config_file}" else: raise ValueError("Either --yaml-config or --template must be provided") + logger_config = None with open(config_file, "r") as fp: - config = replace_env_vars(yaml.safe_load(fp)) + config_contents = yaml.safe_load(fp) + if isinstance(config_contents, dict) and (cfg := config_contents.get("logging_config")): + logger_config = LoggingConfig(**cfg) + logger = get_logger(name=__name__, category="server", config=logger_config) + if args.env: + for env_pair in args.env: + try: + key, value = validate_env_pair(env_pair) + logger.info(f"Setting CLI environment variable {key} => {value}") + os.environ[key] = value + except ValueError as e: + logger.error(f"Error: {str(e)}") + sys.exit(1) + config = replace_env_vars(config_contents) config = StackRunConfig(**config) + # now that the logger is initialized, print the line about which type of config we are using. + logger.info(log_line) + logger.info("Run configuration:") safe_config = redact_sensitive_fields(config.model_dump()) logger.info(yaml.dump(safe_config, indent=2)) diff --git a/llama_stack/log.py b/llama_stack/log.py index 572dea234..0ba95d547 100644 --- a/llama_stack/log.py +++ b/llama_stack/log.py @@ -7,13 +7,15 @@ import logging import os from logging.config import dictConfig -from typing import Dict +from typing import Dict, Optional from rich.console import Console from rich.errors import MarkupError from rich.logging import RichHandler from termcolor import cprint +from .distribution.datatypes import LoggingConfig + # Default log level DEFAULT_LOG_LEVEL = logging.INFO @@ -34,6 +36,56 @@ CATEGORIES = [ _category_levels: Dict[str, int] = {category: DEFAULT_LOG_LEVEL for category in CATEGORIES} +def config_to_category_levels(category: str, level: str): + """ + Helper function to be called either by environment parsing or yaml parsing to go from a list of categories and levels to a dictionary ready to be + used by the logger dictConfig. + + Parameters: + category (str): logging category to apply the level to + level (str): logging level to be used in the category + + Returns: + Dict[str, int]: A dictionary mapping categories to their log levels. + """ + + category_levels: Dict[str, int] = {} + level_value = logging._nameToLevel.get(str(level).upper()) + if level_value is None: + logging.warning(f"Unknown log level '{level}' for category '{category}'. Falling back to default 'INFO'.") + return category_levels + + if category == "all": + # Apply the log level to all categories and the root logger + for cat in CATEGORIES: + category_levels[cat] = level_value + # Set the root logger's level to the specified level + category_levels["root"] = level_value + elif category in CATEGORIES: + category_levels[category] = level_value + logging.info(f"Setting '{category}' category to level '{level}'.") + else: + logging.warning(f"Unknown logging category: {category}. No changes made.") + return category_levels + + +def parse_yaml_config(yaml_config: LoggingConfig) -> Dict[str, int]: + """ + Helper function to parse a yaml logging configuration found in the run.yaml + + Parameters: + yaml_config (Logging): the logger config object found in the run.yaml + + Returns: + Dict[str, int]: A dictionary mapping categories to their log levels. + """ + category_levels = {} + for category, level in yaml_config.category_levels.items(): + category_levels.update(config_to_category_levels(category=category, level=level)) + + return category_levels + + def parse_environment_config(env_config: str) -> Dict[str, int]: """ Parse the LLAMA_STACK_LOGGING environment variable and return a dictionary of category log levels. @@ -53,25 +105,7 @@ def parse_environment_config(env_config: str) -> Dict[str, int]: category, level = pair.split("=", 1) category = category.strip().lower() level = level.strip().upper() # Convert to uppercase for logging._nameToLevel - - level_value = logging._nameToLevel.get(level) - if level_value is None: - logging.warning( - f"Unknown log level '{level}' for category '{category}'. Falling back to default 'INFO'." - ) - continue - - if category == "all": - # Apply the log level to all categories and the root logger - for cat in CATEGORIES: - category_levels[cat] = level_value - # Set the root logger's level to the specified level - category_levels["root"] = level_value - elif category in CATEGORIES: - category_levels[category] = level_value - logging.info(f"Setting '{category}' category to level '{level}'.") - else: - logging.warning(f"Unknown logging category: {category}. No changes made.") + category_levels.update(config_to_category_levels(category=category, level=level)) except ValueError: logging.warning(f"Invalid logging configuration: '{pair}'. Expected format: 'category=level'.") @@ -176,7 +210,9 @@ def setup_logging(category_levels: Dict[str, int], log_file: str | None) -> None logger.setLevel(root_level) -def get_logger(name: str, category: str = "uncategorized") -> logging.LoggerAdapter: +def get_logger( + name: str, category: str = "uncategorized", config: Optional[LoggingConfig] | None = None +) -> logging.LoggerAdapter: """ Returns a logger with the specified name and category. If no category is provided, defaults to 'uncategorized'. @@ -184,10 +220,14 @@ def get_logger(name: str, category: str = "uncategorized") -> logging.LoggerAdap Parameters: name (str): The name of the logger (e.g., module or filename). category (str): The category of the logger (default 'uncategorized'). + config (Logging): optional yaml config to override the existing logger configuration Returns: logging.LoggerAdapter: Configured logger with category support. """ + if config: + _category_levels.update(parse_yaml_config(config)) + logger = logging.getLogger(name) logger.setLevel(_category_levels.get(category, DEFAULT_LOG_LEVEL)) return logging.LoggerAdapter(logger, {"category": category})