Merge 99a3bb1f4e into 09abdb0a37

.pre-commit-config.yaml

@@ -154,6 +154,24 @@ repos:
         pass_filenames: false
         require_serial: true
         files: ^.github/workflows/.*$
+      - id: check-logger-usage
+        name: Check for proper logger usage (use llama_stack.log instead)
+        entry: bash
+        language: system
+        types: [python]
+        pass_filenames: true
+        args:
+          - -c
+          - |
+            matches=$(grep -EnH '^[^#]*\b(import logging|from logging\b)' "$@" | grep -v '# allow-direct-logging' || true)
+            if [ -n "$matches" ]; then
+              # GitHub Actions annotation format
+              while IFS=: read -r file line_num rest; do
+                echo "::error file=$file,line=$line_num::Do not use 'import logging' or 'from logging import' in $file. Use the custom logger instead: from llama_stack.log import get_logger; logger = get_logger(name=__name__, category=\"core\"). If direct logging is truly needed, add: # allow-direct-logging"
+              done <<< "$matches"
+              exit 1
+            fi
+            exit 0
 
 ci:
     autofix_commit_msg: 🎨 [pre-commit.ci] Auto format from pre-commit.com hooks

llama_stack/distribution/build.py

@@ -5,7 +5,6 @@
 # the root directory of this source tree.
 
 import importlib.resources
-import logging
 import sys
 from pathlib import Path
 
@@ -17,11 +16,12 @@ from llama_stack.distribution.distribution import get_provider_registry
 from llama_stack.distribution.external import load_external_apis
 from llama_stack.distribution.utils.exec import run_command
 from llama_stack.distribution.utils.image_types import LlamaStackImageType
+from llama_stack.log import get_logger
+
+logger = get_logger(name=__name__, category="core")
 from llama_stack.providers.datatypes import Api
 from llama_stack.templates.template import DistributionTemplate
 
-log = logging.getLogger(__name__)
-
 # These are the dependencies needed by the distribution server.
 # `llama-stack` is automatically installed by the installation script.
 SERVER_DEPENDENCIES = [
@@ -170,7 +170,7 @@ def build_image(
     return_code = run_command(args)
 
     if return_code != 0:
-        log.error(
+        logger.error(
             f"Failed to build target {image_name} with return code {return_code}",
         )
 

llama_stack/distribution/configure.py

@@ -3,7 +3,6 @@
 #
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
-import logging
 import textwrap
 from typing import Any
 
@@ -21,9 +20,10 @@ from llama_stack.distribution.stack import cast_image_name_to_string, replace_en
 from llama_stack.distribution.utils.config_dirs import EXTERNAL_PROVIDERS_DIR
 from llama_stack.distribution.utils.dynamic import instantiate_class_type
 from llama_stack.distribution.utils.prompt_for_config import prompt_for_config
+from llama_stack.log import get_logger
 from llama_stack.providers.datatypes import Api, ProviderSpec
 
-logger = logging.getLogger(__name__)
+logger = get_logger(name=__name__, category="core")
 
 
 def configure_single_provider(registry: dict[str, ProviderSpec], provider: Provider) -> Provider:

llama_stack/distribution/external.py

@@ -11,7 +11,7 @@ from llama_stack.apis.datatypes import Api, ExternalApiSpec
 from llama_stack.distribution.datatypes import BuildConfig, StackRunConfig
 from llama_stack.log import get_logger
 
-logger = get_logger(name=__name__, category="core")
+logger = get_logger(name=__name__, category="distribution")
 
 
 def load_external_apis(config: StackRunConfig | BuildConfig | None) -> dict[Api, ExternalApiSpec]:

llama_stack/distribution/library_client.py

@@ -7,7 +7,6 @@
 import asyncio
 import inspect
 import json
-import logging
 import os
 import sys
 from concurrent.futures import ThreadPoolExecutor
@@ -32,6 +31,9 @@ from rich.console import Console
 from termcolor import cprint
 
 from llama_stack.distribution.build import print_pip_install_help
+from llama_stack.log import get_logger
+
+logger = get_logger(name=__name__, category="distribution")
 from llama_stack.distribution.configure import parse_and_maybe_upgrade_config
 from llama_stack.distribution.datatypes import Api, BuildConfig, BuildProvider, DistributionSpec
 from llama_stack.distribution.request_headers import (
@@ -48,6 +50,7 @@ from llama_stack.distribution.stack import (
 from llama_stack.distribution.utils.config import redact_sensitive_fields
 from llama_stack.distribution.utils.context import preserve_contexts_async_generator
 from llama_stack.distribution.utils.exec import in_notebook
+from llama_stack.log import get_logger
 from llama_stack.providers.utils.telemetry.tracing import (
     CURRENT_TRACE_CONTEXT,
     end_trace,
@@ -55,7 +58,7 @@ from llama_stack.providers.utils.telemetry.tracing import (
     start_trace,
 )
 
-logger = logging.getLogger(__name__)
+logger = get_logger(name=__name__, category="distribution")
 
 T = TypeVar("T")
 
@@ -173,6 +176,8 @@ class LlamaStackAsLibraryClient(LlamaStackClient):
         """
         Remove all handlers from the root logger. Needed to avoid polluting the console with logs.
         """
+        import logging  # allow-direct-logging
+
         root_logger = logging.getLogger()
 
         for handler in root_logger.handlers[:]:

llama_stack/distribution/providers.py

@@ -16,7 +16,7 @@ from llama_stack.providers.datatypes import HealthResponse, HealthStatus
 from .datatypes import StackRunConfig
 from .utils.config import redact_sensitive_fields
 
-logger = get_logger(name=__name__, category="core")
+logger = get_logger(name=__name__, category="distribution")
 
 
 class ProviderImplConfig(BaseModel):

llama_stack/distribution/request_headers.py

@@ -6,16 +6,16 @@
 
 import contextvars
 import json
-import logging
 from contextlib import AbstractContextManager
 from typing import Any
 
 from llama_stack.distribution.datatypes import User
+from llama_stack.log import get_logger
+
+logger = get_logger(name=__name__, category="distribution")
 
 from .utils.dynamic import instantiate_class_type
 
-log = logging.getLogger(__name__)
-
 # Context variable for request provider data and auth attributes
 PROVIDER_DATA_VAR = contextvars.ContextVar("provider_data", default=None)
 
@@ -61,7 +61,7 @@ class NeedsRequestProviderData:
             provider_data = validator(**val)
             return provider_data
         except Exception as e:
-            log.error(f"Error parsing provider data: {e}")
+            logger.error(f"Error parsing provider data: {e}")
             return None
 
 
@@ -83,7 +83,7 @@ def parse_request_provider_data(headers: dict[str, str]) -> dict[str, Any] | Non
     try:
         return json.loads(val)
     except json.JSONDecodeError:
-        log.error("Provider data not encoded as a JSON object!")
+        logger.error("Provider data not encoded as a JSON object!")
         return None
 
 

llama_stack/distribution/resolver.py

@@ -54,7 +54,7 @@ from llama_stack.providers.datatypes import (
     VectorDBsProtocolPrivate,
 )
 
-logger = get_logger(name=__name__, category="core")
+logger = get_logger(name=__name__, category="distribution")
 
 
 class InvalidProviderError(Exception):

llama_stack/distribution/routers/datasets.py

@@ -12,7 +12,7 @@ from llama_stack.apis.datasets import DatasetPurpose, DataSource
 from llama_stack.log import get_logger
 from llama_stack.providers.datatypes import RoutingTable
 
-logger = get_logger(name=__name__, category="core")
+logger = get_logger(name=__name__, category="distribution")
 
 
 class DatasetIORouter(DatasetIO):

llama_stack/distribution/routers/eval_scoring.py

@@ -16,7 +16,7 @@ from llama_stack.apis.scoring import (
 from llama_stack.log import get_logger
 from llama_stack.providers.datatypes import RoutingTable
 
-logger = get_logger(name=__name__, category="core")
+logger = get_logger(name=__name__, category="distribution")
 
 
 class ScoringRouter(Scoring):

llama_stack/distribution/routers/inference.py

@@ -57,7 +57,7 @@ from llama_stack.providers.utils.inference.inference_store import InferenceStore
 from llama_stack.providers.utils.inference.stream_utils import stream_and_store_openai_completion
 from llama_stack.providers.utils.telemetry.tracing import get_current_span
 
-logger = get_logger(name=__name__, category="core")
+logger = get_logger(name=__name__, category="distribution")
 
 
 class InferenceRouter(Inference):

llama_stack/distribution/routers/safety.py

@@ -14,7 +14,7 @@ from llama_stack.apis.shields import Shield
 from llama_stack.log import get_logger
 from llama_stack.providers.datatypes import RoutingTable
 
-logger = get_logger(name=__name__, category="core")
+logger = get_logger(name=__name__, category="distribution")
 
 
 class SafetyRouter(Safety):

llama_stack/distribution/routers/tool_runtime.py

@@ -22,7 +22,7 @@ from llama_stack.log import get_logger
 
 from ..routing_tables.toolgroups import ToolGroupsRoutingTable
 
-logger = get_logger(name=__name__, category="core")
+logger = get_logger(name=__name__, category="distribution")
 
 
 class ToolRuntimeRouter(ToolRuntime):

llama_stack/distribution/routers/vector_io.py

@@ -30,7 +30,7 @@ from llama_stack.apis.vector_io import (
 from llama_stack.log import get_logger
 from llama_stack.providers.datatypes import HealthResponse, HealthStatus, RoutingTable
 
-logger = get_logger(name=__name__, category="core")
+logger = get_logger(name=__name__, category="distribution")
 
 
 class VectorIORouter(VectorIO):

llama_stack/distribution/routing_tables/benchmarks.py

@@ -14,7 +14,7 @@ from llama_stack.log import get_logger
 
 from .common import CommonRoutingTableImpl
 
-logger = get_logger(name=__name__, category="core")
+logger = get_logger(name=__name__, category="distribution")
 
 
 class BenchmarksRoutingTable(CommonRoutingTableImpl, Benchmarks):

llama_stack/distribution/routing_tables/common.py

@@ -22,7 +22,7 @@ from llama_stack.distribution.store import DistributionRegistry
 from llama_stack.log import get_logger
 from llama_stack.providers.datatypes import Api, RoutingTable
 
-logger = get_logger(name=__name__, category="core")
+logger = get_logger(name=__name__, category="distribution")
 
 
 def get_impl_api(p: Any) -> Api:

llama_stack/distribution/routing_tables/datasets.py

@@ -25,7 +25,7 @@ from llama_stack.log import get_logger
 
 from .common import CommonRoutingTableImpl
 
-logger = get_logger(name=__name__, category="core")
+logger = get_logger(name=__name__, category="distribution")
 
 
 class DatasetsRoutingTable(CommonRoutingTableImpl, Datasets):

llama_stack/distribution/routing_tables/models.py

@@ -16,7 +16,7 @@ from llama_stack.log import get_logger
 
 from .common import CommonRoutingTableImpl, lookup_model
 
-logger = get_logger(name=__name__, category="core")
+logger = get_logger(name=__name__, category="distribution")
 
 
 class ModelsRoutingTable(CommonRoutingTableImpl, Models):

llama_stack/distribution/routing_tables/scoring_functions.py

@@ -19,7 +19,7 @@ from llama_stack.log import get_logger
 
 from .common import CommonRoutingTableImpl
 
-logger = get_logger(name=__name__, category="core")
+logger = get_logger(name=__name__, category="distribution")
 
 
 class ScoringFunctionsRoutingTable(CommonRoutingTableImpl, ScoringFunctions):

llama_stack/distribution/routing_tables/shields.py

@@ -15,7 +15,7 @@ from llama_stack.log import get_logger
 
 from .common import CommonRoutingTableImpl
 
-logger = get_logger(name=__name__, category="core")
+logger = get_logger(name=__name__, category="distribution")
 
 
 class ShieldsRoutingTable(CommonRoutingTableImpl, Shields):

llama_stack/distribution/routing_tables/toolgroups.py

@@ -13,7 +13,7 @@ from llama_stack.log import get_logger
 
 from .common import CommonRoutingTableImpl
 
-logger = get_logger(name=__name__, category="core")
+logger = get_logger(name=__name__, category="distribution")
 
 
 def parse_toolgroup_from_toolgroup_name_pair(toolgroup_name_with_maybe_tool_name: str) -> str | None:

llama_stack/distribution/routing_tables/vector_dbs.py

@@ -29,7 +29,7 @@ from llama_stack.log import get_logger
 
 from .common import CommonRoutingTableImpl, lookup_model
 
-logger = get_logger(name=__name__, category="core")
+logger = get_logger(name=__name__, category="distribution")
 
 
 class VectorDBsRoutingTable(CommonRoutingTableImpl, VectorDBs):

llama_stack/distribution/server/server.py

@@ -9,7 +9,7 @@ import asyncio
 import functools
 import inspect
 import json
-import logging
+import logging  # allow-direct-logging
 import os
 import ssl
 import sys

llama_stack/distribution/stack.py

@@ -45,7 +45,7 @@ from llama_stack.distribution.utils.dynamic import instantiate_class_type
 from llama_stack.log import get_logger
 from llama_stack.providers.datatypes import Api
 
-logger = get_logger(name=__name__, category="core")
+logger = get_logger(name=__name__, category="distribution")
 
 
 class LlamaStack(

llama_stack/distribution/utils/exec.py

@@ -4,21 +4,20 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-import logging
+import importlib
+import json
 import os
 import signal
 import subprocess
 import sys
+from pathlib import Path
 
 from termcolor import cprint
 
-log = logging.getLogger(__name__)
-
-import importlib
-import json
-from pathlib import Path
-
 from llama_stack.distribution.utils.image_types import LlamaStackImageType
+from llama_stack.log import get_logger
+
+logger = get_logger(name=__name__, category="distribution")
 
 
 def formulate_run_args(image_type: str, image_name: str) -> list[str]:
@@ -119,7 +118,7 @@ def run_command(command: list[str]) -> int:
     def sigint_handler(signum, frame):
         nonlocal ctrl_c_pressed
         ctrl_c_pressed = True
-        log.info("\nCtrl-C detected. Aborting...")
+        logger.info("\nCtrl-C detected. Aborting...")
 
     try:
         # Set up the signal handler
@@ -133,10 +132,10 @@ def run_command(command: list[str]) -> int:
         )
         return result.returncode
     except subprocess.SubprocessError as e:
-        log.error(f"Subprocess error: {e}")
+        logger.error(f"Subprocess error: {e}")
         return 1
     except Exception as e:
-        log.exception(f"Unexpected error: {e}")
+        logger.exception(f"Unexpected error: {e}")
         return 1
     finally:
         # Restore the original signal handler

llama_stack/distribution/utils/prompt_for_config.py

@@ -6,7 +6,6 @@
 
 import inspect
 import json
-import logging
 from enum import Enum
 from typing import Annotated, Any, Literal, Union, get_args, get_origin
 
@@ -14,7 +13,9 @@ from pydantic import BaseModel
 from pydantic.fields import FieldInfo
 from pydantic_core import PydanticUndefinedType
 
-log = logging.getLogger(__name__)
+from llama_stack.log import get_logger
+
+logger = get_logger(name=__name__, category="distribution")
 
 
 def is_list_of_primitives(field_type):
@@ -107,7 +108,7 @@ def prompt_for_discriminated_union(
 
         if discriminator_value in type_map:
             chosen_type = type_map[discriminator_value]
-            log.info(f"\nConfiguring {chosen_type.__name__}:")
+            logger.info(f"\nConfiguring {chosen_type.__name__}:")
 
             if existing_value and (getattr(existing_value, discriminator) != discriminator_value):
                 existing_value = None
@@ -117,7 +118,7 @@ def prompt_for_discriminated_union(
             setattr(sub_config, discriminator, discriminator_value)
             return sub_config
         else:
-            log.error(f"Invalid {discriminator}. Please try again.")
+            logger.error(f"Invalid {discriminator}. Please try again.")
 
 
 # This is somewhat elaborate, but does not purport to be comprehensive in any way.
@@ -166,7 +167,7 @@ def prompt_for_config(config_type: type[BaseModel], existing_config: BaseModel |
                     config_data[field_name] = validated_value
                     break
                 except KeyError:
-                    log.error(f"Invalid choice. Please choose from: {', '.join(e.name for e in field_type)}")
+                    logger.error(f"Invalid choice. Please choose from: {', '.join(e.name for e in field_type)}")
             continue
 
         if is_discriminated_union(field):
@@ -179,7 +180,7 @@ def prompt_for_config(config_type: type[BaseModel], existing_config: BaseModel |
                 config_data[field_name] = None
                 continue
             nested_type = get_non_none_type(field_type)
-            log.info(f"Entering sub-configuration for {field_name}:")
+            logger.info(f"Entering sub-configuration for {field_name}:")
             config_data[field_name] = prompt_for_config(nested_type, existing_value)
         elif is_optional(field_type) and is_discriminated_union(get_non_none_type(field_type)):
             prompt = f"Do you want to configure {field_name}? (y/n): "
@@ -193,7 +194,7 @@ def prompt_for_config(config_type: type[BaseModel], existing_config: BaseModel |
                 existing_value,
             )
         elif can_recurse(field_type):
-            log.info(f"\nEntering sub-configuration for {field_name}:")
+            logger.info(f"\nEntering sub-configuration for {field_name}:")
             config_data[field_name] = prompt_for_config(
                 field_type,
                 existing_value,
@@ -220,7 +221,7 @@ def prompt_for_config(config_type: type[BaseModel], existing_config: BaseModel |
                         config_data[field_name] = None
                         break
                     else:
-                        log.error("This field is required. Please provide a value.")
+                        logger.error("This field is required. Please provide a value.")
                         continue
                 else:
                     try:
@@ -242,10 +243,10 @@ def prompt_for_config(config_type: type[BaseModel], existing_config: BaseModel |
                                 value = [element_type(item) for item in value]
 
                             except json.JSONDecodeError:
-                                log.error('Invalid JSON. Please enter a valid JSON-encoded list e.g., ["foo","bar"]')
+                                logger.error('Invalid JSON. Please enter a valid JSON-encoded list e.g., ["foo","bar"]')
                                 continue
                             except ValueError as e:
-                                log.error(f"{str(e)}")
+                                logger.error(f"{str(e)}")
                                 continue
 
                         elif get_origin(field_type) is dict:
@@ -255,7 +256,7 @@ def prompt_for_config(config_type: type[BaseModel], existing_config: BaseModel |
                                     raise ValueError("Input must be a JSON-encoded dictionary")
 
                             except json.JSONDecodeError:
-                                log.error("Invalid JSON. Please enter a valid JSON-encoded dict.")
+                                logger.error("Invalid JSON. Please enter a valid JSON-encoded dict.")
                                 continue
 
                         # Convert the input to the correct type
@@ -268,7 +269,9 @@ def prompt_for_config(config_type: type[BaseModel], existing_config: BaseModel |
                             value = field_type(user_input)
 
                     except ValueError:
-                        log.error(f"Invalid input. Expected type: {getattr(field_type, '__name__', str(field_type))}")
+                        logger.error(
+                            f"Invalid input. Expected type: {getattr(field_type, '__name__', str(field_type))}"
+                        )
                         continue
 
                 try:
@@ -277,6 +280,6 @@ def prompt_for_config(config_type: type[BaseModel], existing_config: BaseModel |
                     config_data[field_name] = validated_value
                     break
                 except ValueError as e:
-                    log.error(f"Validation error: {str(e)}")
+                    logger.error(f"Validation error: {str(e)}")
 
     return config_type(**config_data)

llama_stack/log.py

@@ -4,11 +4,11 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-import logging
+import logging  # allow-direct-logging
 import os
 import re
 import sys
-from logging.config import dictConfig
+from logging.config import dictConfig  # allow-direct-logging
 
 from rich.console import Console
 from rich.errors import MarkupError

llama_stack/models/llama/llama3/multimodal/encoder_utils.py

@@ -13,7 +13,7 @@
 
 # Copyright (c) Meta Platforms, Inc. and its affiliates.
 import math
-from logging import getLogger
+from logging import getLogger  # allow-direct-logging
 
 import torch
 import torch.nn.functional as F

llama_stack/models/llama/llama3/multimodal/image_transform.py

@@ -13,7 +13,7 @@
 
 import math
 from collections import defaultdict
-from logging import getLogger
+from logging import getLogger  # allow-direct-logging
 from typing import Any
 
 import torch

llama_stack/models/llama/llama3/multimodal/model.py

@@ -4,7 +4,6 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-import logging
 import math
 from collections.abc import Callable
 from functools import partial
@@ -22,6 +21,8 @@ from PIL import Image as PIL_Image
 from torch import Tensor, nn
 from torch.distributed import _functional_collectives as funcol
 
+from llama_stack.log import get_logger
+
 from ..model import ModelArgs, RMSNorm, apply_rotary_emb, precompute_freqs_cis
 from .encoder_utils import (
     build_encoder_attention_mask,
@@ -34,9 +35,10 @@ from .encoder_utils import (
 from .image_transform import VariableSizeImageTransform
 from .utils import get_negative_inf_value, to_2tuple
 
-logger = logging.getLogger(__name__)
 MP_SCALE = 8
 
+logger = get_logger(name=__name__, category="core")
+
 
 def reduce_from_tensor_model_parallel_region(input_):
     """All-reduce the input tensor across model parallel group."""
@@ -771,7 +773,7 @@ class TilePositionEmbedding(nn.Module):
         if embed is not None:
             # reshape the weights to the correct shape
             nt_old, nt_old, _, w = embed.shape
-            logging.info(f"Resizing tile embedding from {nt_old}x{nt_old} to {self.num_tiles}x{self.num_tiles}")
+            logger.info(f"Resizing tile embedding from {nt_old}x{nt_old} to {self.num_tiles}x{self.num_tiles}")
             embed_new = TilePositionEmbedding._dynamic_resize(embed, self.num_tiles)
             # assign the weights to the module
             state_dict[prefix + "embedding"] = embed_new

llama_stack/models/llama/llama3/tokenizer.py

@@ -5,7 +5,7 @@
 # the root directory of this source tree.
 
 from collections.abc import Collection, Iterator, Sequence, Set
-from logging import getLogger
+from logging import getLogger  # allow-direct-logging
 from pathlib import Path
 from typing import (
     Literal,

llama_stack/models/llama/llama4/quantization/loader.py

@@ -4,7 +4,6 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-import logging
 import os
 from collections.abc import Callable
 
@@ -13,12 +12,14 @@ from fairscale.nn.model_parallel.initialize import get_model_parallel_rank
 from torch import Tensor, nn
 from torch.nn import functional as F
 
+from llama_stack.log import get_logger
+
+log = get_logger(__name__, category="core")
+
 from ...datatypes import QuantizationMode
 from ..model import Transformer, TransformerBlock
 from ..moe import MoE
 
-log = logging.getLogger(__name__)
-
 
 def swiglu_wrapper_no_reduce(
     self,

llama_stack/models/llama/llama4/tokenizer.py

@@ -5,7 +5,7 @@
 # the root directory of this source tree.
 
 from collections.abc import Collection, Iterator, Sequence, Set
-from logging import getLogger
+from logging import getLogger  # allow-direct-logging
 from pathlib import Path
 from typing import (
     Literal,

llama_stack/models/llama/quantize_impls.py

@@ -6,9 +6,10 @@
 
 # type: ignore
 import collections
-import logging
 
-log = logging.getLogger(__name__)
+from llama_stack.log import get_logger
+
+log = get_logger(__name__, category="core")
 
 try:
     import fbgemm_gpu.experimental.gen_ai  # noqa: F401

llama_stack/providers/inline/agents/meta_reference/agents.py

@@ -4,7 +4,6 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-import logging
 import uuid
 from collections.abc import AsyncGenerator
 from datetime import UTC, datetime
@@ -42,6 +41,7 @@ from llama_stack.apis.safety import Safety
 from llama_stack.apis.tools import ToolGroups, ToolRuntime
 from llama_stack.apis.vector_io import VectorIO
 from llama_stack.distribution.datatypes import AccessRule
+from llama_stack.log import get_logger
 from llama_stack.providers.utils.kvstore import InmemoryKVStoreImpl, kvstore_impl
 from llama_stack.providers.utils.pagination import paginate_records
 from llama_stack.providers.utils.responses.responses_store import ResponsesStore
@@ -51,7 +51,7 @@ from .config import MetaReferenceAgentsImplConfig
 from .openai_responses import OpenAIResponsesImpl
 from .persistence import AgentInfo
 
-logger = logging.getLogger()
+logger = get_logger(name=__name__, category="agents")
 
 
 class MetaReferenceAgentsImpl(Agents):

llama_stack/providers/inline/agents/meta_reference/persistence.py

@@ -5,7 +5,6 @@
 # the root directory of this source tree.
 
 import json
-import logging
 import uuid
 from datetime import UTC, datetime
 
@@ -14,9 +13,10 @@ from llama_stack.distribution.access_control.access_control import AccessDeniedE
 from llama_stack.distribution.access_control.datatypes import AccessRule
 from llama_stack.distribution.datatypes import User
 from llama_stack.distribution.request_headers import get_authenticated_user
-from llama_stack.providers.utils.kvstore import KVStore
+from llama_stack.log import get_logger
 
-log = logging.getLogger(__name__)
+logger = get_logger(name=__name__, category="agents")
+from llama_stack.providers.utils.kvstore import KVStore
 
 
 class AgentSessionInfo(Session):
@@ -126,7 +126,7 @@ class AgentPersistence:
                 turn = Turn(**json.loads(value))
                 turns.append(turn)
             except Exception as e:
-                log.error(f"Error parsing turn: {e}")
+                logger.error(f"Error parsing turn: {e}")
                 continue
 
         # The kvstore does not guarantee order, so we sort by started_at
@@ -193,7 +193,7 @@ class AgentPersistence:
                 session_info = Session(**json.loads(value))
                 sessions.append(session_info)
             except Exception as e:
-                log.error(f"Error parsing session info: {e}")
+                logger.error(f"Error parsing session info: {e}")
                 continue
         return sessions
 

llama_stack/providers/inline/agents/meta_reference/safety.py

@@ -5,13 +5,13 @@
 # the root directory of this source tree.
 
 import asyncio
-import logging
 
 from llama_stack.apis.inference import Message
 from llama_stack.apis.safety import Safety, SafetyViolation, ViolationLevel
-from llama_stack.providers.utils.telemetry import tracing
+from llama_stack.log import get_logger
 
-log = logging.getLogger(__name__)
+logger = get_logger(name=__name__, category="agents")
+from llama_stack.providers.utils.telemetry import tracing
 
 
 class SafetyException(Exception):  # noqa: N818
@@ -49,4 +49,4 @@ class ShieldRunnerMixin:
             if violation.violation_level == ViolationLevel.ERROR:
                 raise SafetyException(violation)
             elif violation.violation_level == ViolationLevel.WARN:
-                log.warning(f"[Warn]{identifier} raised a warning")
+                logger.warning(f"[Warn]{identifier} raised a warning")

llama_stack/providers/inline/inference/meta_reference/parallel_utils.py

@@ -12,7 +12,6 @@
 
 import copy
 import json
-import logging
 import multiprocessing
 import os
 import tempfile
@@ -32,14 +31,15 @@ from fairscale.nn.model_parallel.initialize import (
 from pydantic import BaseModel, Field
 from torch.distributed.launcher.api import LaunchConfig, elastic_launch
 
+from llama_stack.log import get_logger
+
+logger = get_logger(name=__name__, category="inference")
 from llama_stack.models.llama.datatypes import GenerationResult
 from llama_stack.providers.utils.inference.prompt_adapter import (
     ChatCompletionRequestWithRawContent,
     CompletionRequestWithRawContent,
 )
 
-log = logging.getLogger(__name__)
-
 
 class ProcessingMessageName(str, Enum):
     ready_request = "ready_request"
@@ -169,13 +169,13 @@ def retrieve_requests(reply_socket_url: str):
                         group=get_model_parallel_group(),
                     )
                     if isinstance(updates[0], CancelSentinel):
-                        log.info("quitting generation loop because request was cancelled")
+                        logger.info("quitting generation loop because request was cancelled")
                         break
 
                 if mp_rank_0():
                     send_obj(EndSentinel())
             except Exception as e:
-                log.exception("exception in generation loop")
+                logger.exception("exception in generation loop")
 
                 if mp_rank_0():
                     send_obj(ExceptionResponse(error=str(e)))
@@ -236,7 +236,7 @@ def worker_process_entrypoint(
         except StopIteration:
             break
 
-    log.info("[debug] worker process done")
+            logger.info("[debug] worker process done")
 
 
 def launch_dist_group(
@@ -294,7 +294,7 @@ def start_model_parallel_process(
 
     request_socket.send(encode_msg(ReadyRequest()))
     _response = request_socket.recv()
-    log.info("Loaded model...")
+    logger.info("Loaded model...")
 
     return request_socket, process
 
@@ -346,7 +346,7 @@ class ModelParallelProcessGroup:
                     break
 
                 if isinstance(obj, ExceptionResponse):
-                    log.error(f"[debug] got exception {obj.error}")
+                    logger.error(f"[debug] got exception {obj.error}")
                     raise Exception(obj.error)
 
                 if isinstance(obj, TaskResponse):

llama_stack/providers/inline/inference/sentence_transformers/sentence_transformers.py

@@ -4,7 +4,6 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-import logging
 from collections.abc import AsyncGenerator
 
 from llama_stack.apis.inference import (
@@ -32,8 +31,6 @@ from llama_stack.providers.utils.inference.openai_compat import (
 
 from .config import SentenceTransformersInferenceConfig
 
-log = logging.getLogger(__name__)
-
 
 class SentenceTransformersInferenceImpl(
     OpenAIChatCompletionToLlamaStackMixin,

llama_stack/providers/inline/post_training/huggingface/recipes/finetune_single_device.py

@@ -6,7 +6,6 @@
 
 import gc
 import json
-import logging
 import multiprocessing
 import os
 import signal
@@ -45,10 +44,11 @@ from llama_stack.apis.post_training import (
     LoraFinetuningConfig,
     TrainingConfig,
 )
+from llama_stack.log import get_logger
 
 from ..config import HuggingFacePostTrainingConfig
 
-logger = logging.getLogger(__name__)
+logger = get_logger(name=__name__, category="core")
 
 
 def get_gb(to_convert: int) -> str:

llama_stack/providers/inline/post_training/torchtune/recipes/lora_finetuning_single_device.py

@@ -4,7 +4,6 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-import logging
 import os
 import time
 from datetime import UTC, datetime
@@ -19,6 +18,7 @@ from torch.utils.data import DataLoader, DistributedSampler
 from torchtune import modules, training
 from torchtune import utils as torchtune_utils
 from torchtune.data import padded_collate_sft
+from torchtune.models.llama3._tokenizer import Llama3Tokenizer
 from torchtune.modules.loss import CEWithChunkedOutputLoss
 from torchtune.modules.peft import (
     get_adapter_params,
@@ -45,6 +45,9 @@ from llama_stack.apis.post_training import (
 )
 from llama_stack.distribution.utils.config_dirs import DEFAULT_CHECKPOINT_DIR
 from llama_stack.distribution.utils.model_utils import model_local_dir
+from llama_stack.log import get_logger
+
+logger = get_logger(name=__name__, category="core")
 from llama_stack.models.llama.sku_list import resolve_model
 from llama_stack.providers.inline.post_training.common.utils import evacuate_model_from_device
 from llama_stack.providers.inline.post_training.torchtune.common import utils
@@ -56,10 +59,6 @@ from llama_stack.providers.inline.post_training.torchtune.config import (
 )
 from llama_stack.providers.inline.post_training.torchtune.datasets.sft import SFTDataset
 
-log = logging.getLogger(__name__)
-
-from torchtune.models.llama3._tokenizer import Llama3Tokenizer
-
 
 class LoraFinetuningSingleDevice:
     # This recipe doesn't include several training efficiency setting within origin torchtune repo, including
@@ -183,18 +182,18 @@ class LoraFinetuningSingleDevice:
             base_model_state_dict=checkpoint_dict[training.MODEL_KEY],
             lora_weights_state_dict=None,
         )
-        log.info(f"Model is initialized with precision {self._dtype}.")
+        logger.info(f"Model is initialized with precision {self._dtype}.")
 
         self._tokenizer = await self._setup_tokenizer()
-        log.info("Tokenizer is initialized.")
+        logger.info("Tokenizer is initialized.")
 
         assert isinstance(self.training_config.optimizer_config, OptimizerConfig), "OptimizerConfig must be initialized"
         self._optimizer = await self._setup_optimizer(optimizer_config=self.training_config.optimizer_config)
-        log.info("Optimizer is initialized.")
+        logger.info("Optimizer is initialized.")
 
         self._loss_fn = CEWithChunkedOutputLoss()
         self._model.set_num_output_chunks(self._loss_fn.num_output_chunks)
-        log.info("Loss is initialized.")
+        logger.info("Loss is initialized.")
 
         assert isinstance(self.training_config.data_config, DataConfig), "DataConfig must be initialized"
 
@@ -213,7 +212,7 @@ class LoraFinetuningSingleDevice:
                 batch_size=self._batch_size,
             )
 
-        log.info("Dataset and Sampler are initialized.")
+        logger.info("Dataset and Sampler are initialized.")
 
         # Number of training steps in each epoch depends on the number of batches produced
         # by the dataloader and the max_steps_per_epoch param set by the user and is used
@@ -231,7 +230,7 @@ class LoraFinetuningSingleDevice:
             num_training_steps=self.total_epochs * self._steps_per_epoch,
             last_epoch=self.global_step - 1,
         )
-        log.info("Learning rate scheduler is initialized.")
+        logger.info("Learning rate scheduler is initialized.")
 
         # Used to ignore labels for loss computation
         self.ignore_labels_cache = torch.full((self._batch_size, 1), self._loss_fn.ignore_index, device=self._device)
@@ -533,7 +532,7 @@ class LoraFinetuningSingleDevice:
                     t0 = time.perf_counter()
 
             self.epochs_run += 1
-            log.info("Starting checkpoint save...")
+            logger.info("Starting checkpoint save...")
             checkpoint_path = await self.save_checkpoint(epoch=curr_epoch)
             checkpoint = Checkpoint(
                 identifier=f"{self.model_id}-sft-{curr_epoch}",
@@ -561,7 +560,7 @@ class LoraFinetuningSingleDevice:
     async def validation(self) -> tuple[float, float]:
         total_loss = 0.0
         total_tokens = 0
-        log.info("Starting validation...")
+        logger.info("Starting validation...")
         pbar = tqdm(total=len(self._validation_dataloader))
         for idx, batch in enumerate(self._validation_dataloader):
             if idx == self.max_validation_steps:

llama_stack/providers/inline/safety/code_scanner/code_scanner.py

@@ -4,7 +4,6 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-import logging
 from typing import Any
 
 from llama_stack.apis.inference import Message
@@ -15,14 +14,15 @@ from llama_stack.apis.safety import (
     ViolationLevel,
 )
 from llama_stack.apis.shields import Shield
+from llama_stack.log import get_logger
+
+logger = get_logger(name=__name__, category="safety")
 from llama_stack.providers.utils.inference.prompt_adapter import (
     interleaved_content_as_str,
 )
 
 from .config import CodeScannerConfig
 
-log = logging.getLogger(__name__)
-
 ALLOWED_CODE_SCANNER_MODEL_IDS = [
     "CodeScanner",
     "CodeShield",
@@ -58,7 +58,7 @@ class MetaReferenceCodeScannerSafetyImpl(Safety):
         from codeshield.cs import CodeShield
 
         text = "\n".join([interleaved_content_as_str(m.content) for m in messages])
-        log.info(f"Running CodeScannerShield on {text[50:]}")
+        logger.info(f"Running CodeScannerShield on {text[50:]}")
         result = await CodeShield.scan_code(text)
 
         violation = None

llama_stack/providers/inline/safety/prompt_guard/prompt_guard.py

@@ -4,7 +4,6 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-import logging
 from typing import Any
 
 import torch
@@ -19,6 +18,9 @@ from llama_stack.apis.safety import (
 )
 from llama_stack.apis.shields import Shield
 from llama_stack.distribution.utils.model_utils import model_local_dir
+from llama_stack.log import get_logger
+
+logger = get_logger(name=__name__, category="safety")
 from llama_stack.providers.datatypes import ShieldsProtocolPrivate
 from llama_stack.providers.utils.inference.prompt_adapter import (
     interleaved_content_as_str,
@@ -26,8 +28,6 @@ from llama_stack.providers.utils.inference.prompt_adapter import (
 
 from .config import PromptGuardConfig, PromptGuardType
 
-log = logging.getLogger(__name__)
-
 PROMPT_GUARD_MODEL = "Prompt-Guard-86M"
 
 
@@ -96,7 +96,7 @@ class PromptGuardShield:
         probabilities = torch.softmax(logits / self.temperature, dim=-1)
         score_embedded = probabilities[0, 1].item()
         score_malicious = probabilities[0, 2].item()
-        log.info(
+        logger.info(
             f"Ran PromptGuardShield and got Scores: Embedded: {score_embedded}, Malicious: {score_malicious}",
         )
 

llama_stack/providers/inline/scoring/basic/utils/ifeval_utils.py

@@ -7,7 +7,6 @@
 import collections
 import functools
 import json
-import logging
 import random
 import re
 import string
@@ -20,7 +19,9 @@ import nltk
 from pythainlp.tokenize import sent_tokenize as sent_tokenize_thai
 from pythainlp.tokenize import word_tokenize as word_tokenize_thai
 
-logger = logging.getLogger()
+from llama_stack.log import get_logger
+
+logger = get_logger(name=__name__, category="core")
 
 WORD_LIST = [
     "western",

llama_stack/providers/inline/tool_runtime/rag/memory.py

@@ -5,7 +5,6 @@
 # the root directory of this source tree.
 
 import asyncio
-import logging
 import secrets
 import string
 from typing import Any
@@ -32,6 +31,9 @@ from llama_stack.apis.tools import (
     ToolRuntime,
 )
 from llama_stack.apis.vector_io import QueryChunksResponse, VectorIO
+from llama_stack.log import get_logger
+
+logger = get_logger(name=__name__, category="tools")
 from llama_stack.providers.datatypes import ToolGroupsProtocolPrivate
 from llama_stack.providers.utils.inference.prompt_adapter import interleaved_content_as_str
 from llama_stack.providers.utils.memory.vector_store import (
@@ -42,8 +44,6 @@ from llama_stack.providers.utils.memory.vector_store import (
 from .config import RagToolRuntimeConfig
 from .context_retriever import generate_rag_query
 
-log = logging.getLogger(__name__)
-
 
 def make_random_string(length: int = 8):
     return "".join(secrets.choice(string.ascii_letters + string.digits) for _ in range(length))
@@ -153,7 +153,7 @@ class MemoryToolRuntimeImpl(ToolGroupsProtocolPrivate, ToolRuntime, RAGToolRunti
             tokens += metadata.get("metadata_token_count", 0)
 
             if tokens > query_config.max_tokens_in_context:
-                log.error(
+                logger.error(
                     f"Using {len(picked)} chunks; reached max tokens in context: {tokens}",
                 )
                 break

llama_stack/providers/inline/vector_io/faiss/faiss.py

@@ -8,7 +8,6 @@ import asyncio
 import base64
 import io
 import json
-import logging
 from typing import Any
 
 import faiss
@@ -23,6 +22,7 @@ from llama_stack.apis.vector_io import (
     QueryChunksResponse,
     VectorIO,
 )
+from llama_stack.log import get_logger
 from llama_stack.providers.datatypes import (
     HealthResponse,
     HealthStatus,
@@ -38,7 +38,7 @@ from llama_stack.providers.utils.memory.vector_store import (
 
 from .config import FaissVectorIOConfig
 
-logger = logging.getLogger(__name__)
+logger = get_logger(name=__name__, category="core")
 
 VERSION = "v3"
 VECTOR_DBS_PREFIX = f"vector_dbs:{VERSION}::"

llama_stack/providers/inline/vector_io/sqlite_vec/sqlite_vec.py

@@ -5,7 +5,6 @@
 # the root directory of this source tree.
 
 import asyncio
-import logging
 import re
 import sqlite3
 import struct
@@ -23,6 +22,7 @@ from llama_stack.apis.vector_io import (
     QueryChunksResponse,
     VectorIO,
 )
+from llama_stack.log import get_logger
 from llama_stack.providers.datatypes import VectorDBsProtocolPrivate
 from llama_stack.providers.utils.kvstore import kvstore_impl
 from llama_stack.providers.utils.kvstore.api import KVStore
@@ -34,7 +34,7 @@ from llama_stack.providers.utils.memory.vector_store import (
     VectorDBWithIndex,
 )
 
-logger = logging.getLogger(__name__)
+logger = get_logger(name=__name__, category="core")
 
 # Specifying search mode is dependent on the VectorIO provider.
 VECTOR_SEARCH = "vector"

llama_stack/providers/remote/inference/llama_openai_compat/llama.py

@@ -3,7 +3,6 @@
 #
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
-import logging
 
 from llama_stack.providers.remote.inference.llama_openai_compat.config import LlamaCompatConfig
 from llama_stack.providers.utils.inference.litellm_openai_mixin import LiteLLMOpenAIMixin
@@ -11,8 +10,6 @@ from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin
 
 from .models import MODEL_ENTRIES
 
-logger = logging.getLogger(__name__)
-
 
 class LlamaCompatInferenceAdapter(OpenAIMixin, LiteLLMOpenAIMixin):
     """

llama_stack/providers/remote/inference/nvidia/nvidia.py

@@ -4,7 +4,6 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-import logging
 import warnings
 from collections.abc import AsyncIterator
 
@@ -33,6 +32,7 @@ from llama_stack.apis.inference import (
     ToolChoice,
     ToolConfig,
 )
+from llama_stack.log import get_logger
 from llama_stack.models.llama.datatypes import ToolDefinition, ToolPromptFormat
 from llama_stack.providers.utils.inference.model_registry import (
     ModelRegistryHelper,
@@ -54,7 +54,7 @@ from .openai_utils import (
 )
 from .utils import _is_nvidia_hosted
 
-logger = logging.getLogger(__name__)
+logger = get_logger(name=__name__, category="inference")
 
 
 class NVIDIAInferenceAdapter(OpenAIMixin, Inference, ModelRegistryHelper):

llama_stack/providers/remote/inference/nvidia/utils.py

@@ -4,13 +4,14 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-import logging
 
 import httpx
 
+from llama_stack.log import get_logger
+
 from . import NVIDIAConfig
 
-logger = logging.getLogger(__name__)
+logger = get_logger(name=__name__, category="inference")
 
 
 def _is_nvidia_hosted(config: NVIDIAConfig) -> bool:

llama_stack/providers/remote/inference/openai/openai.py

@@ -4,7 +4,6 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-import logging
 
 from llama_stack.providers.utils.inference.litellm_openai_mixin import LiteLLMOpenAIMixin
 from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin
@@ -12,8 +11,6 @@ from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin
 from .config import OpenAIConfig
 from .models import MODEL_ENTRIES
 
-logger = logging.getLogger(__name__)
-
 
 #
 # This OpenAI adapter implements Inference methods using two mixins -

llama_stack/providers/remote/inference/tgi/tgi.py

@@ -5,7 +5,6 @@
 # the root directory of this source tree.
 
 
-import logging
 from collections.abc import AsyncGenerator
 
 from huggingface_hub import AsyncInferenceClient, HfApi
@@ -34,6 +33,9 @@ from llama_stack.apis.inference import (
     ToolPromptFormat,
 )
 from llama_stack.apis.models import Model
+from llama_stack.log import get_logger
+
+log = get_logger(__name__, category="core")
 from llama_stack.models.llama.sku_list import all_registered_models
 from llama_stack.providers.datatypes import ModelsProtocolPrivate
 from llama_stack.providers.utils.inference.model_registry import (
@@ -58,8 +60,6 @@ from llama_stack.providers.utils.inference.prompt_adapter import (
 
 from .config import InferenceAPIImplConfig, InferenceEndpointImplConfig, TGIImplConfig
 
-log = logging.getLogger(__name__)
-
 
 def build_hf_repo_model_entries():
     return [

llama_stack/providers/remote/post_training/nvidia/utils.py

@@ -4,7 +4,6 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-import logging
 import warnings
 from typing import Any
 
@@ -15,8 +14,6 @@ from llama_stack.providers.remote.post_training.nvidia.config import SFTLoRADefa
 
 from .config import NvidiaPostTrainingConfig
 
-logger = logging.getLogger(__name__)
-
 
 def warn_unsupported_params(config_dict: Any, supported_keys: set[str], config_name: str) -> None:
     keys = set(config_dict.__annotations__.keys()) if isinstance(config_dict, BaseModel) else config_dict.keys()

llama_stack/providers/remote/safety/bedrock/bedrock.py

@@ -5,7 +5,6 @@
 # the root directory of this source tree.
 
 import json
-import logging
 from typing import Any
 
 from llama_stack.apis.inference import Message
@@ -16,12 +15,13 @@ from llama_stack.apis.safety import (
     ViolationLevel,
 )
 from llama_stack.apis.shields import Shield
+from llama_stack.log import get_logger
 from llama_stack.providers.datatypes import ShieldsProtocolPrivate
 from llama_stack.providers.utils.bedrock.client import create_bedrock_client
 
 from .config import BedrockSafetyConfig
 
-logger = logging.getLogger(__name__)
+logger = get_logger(name=__name__, category="safety")
 
 
 class BedrockSafetyAdapter(Safety, ShieldsProtocolPrivate):

llama_stack/providers/remote/safety/nvidia/nvidia.py

@@ -4,7 +4,6 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-import logging
 from typing import Any
 
 import requests
@@ -17,8 +16,6 @@ from llama_stack.providers.utils.inference.openai_compat import convert_message_
 
 from .config import NVIDIASafetyConfig
 
-logger = logging.getLogger(__name__)
-
 
 class NVIDIASafetyAdapter(Safety, ShieldsProtocolPrivate):
     def __init__(self, config: NVIDIASafetyConfig) -> None:

llama_stack/providers/remote/safety/sambanova/sambanova.py

@@ -5,7 +5,6 @@
 # the root directory of this source tree.
 
 import json
-import logging
 from typing import Any
 
 import litellm
@@ -20,12 +19,13 @@ from llama_stack.apis.safety import (
 )
 from llama_stack.apis.shields import Shield
 from llama_stack.distribution.request_headers import NeedsRequestProviderData
+from llama_stack.log import get_logger
 from llama_stack.providers.datatypes import ShieldsProtocolPrivate
 from llama_stack.providers.utils.inference.openai_compat import convert_message_to_openai_dict_new
 
 from .config import SambaNovaSafetyConfig
 
-logger = logging.getLogger(__name__)
+logger = get_logger(name=__name__, category="safety")
 
 CANNED_RESPONSE_TEXT = "I can't answer that. Can I help with something else?"
 

llama_stack/providers/remote/vector_io/chroma/chroma.py

@@ -5,7 +5,6 @@
 # the root directory of this source tree.
 import asyncio
 import json
-import logging
 from typing import Any
 from urllib.parse import urlparse
 
@@ -20,6 +19,9 @@ from llama_stack.apis.vector_io import (
     QueryChunksResponse,
     VectorIO,
 )
+from llama_stack.log import get_logger
+
+log = get_logger(__name__, category="core")
 from llama_stack.providers.datatypes import Api, VectorDBsProtocolPrivate
 from llama_stack.providers.inline.vector_io.chroma import ChromaVectorIOConfig as InlineChromaVectorIOConfig
 from llama_stack.providers.utils.kvstore import kvstore_impl
@@ -32,8 +34,6 @@ from llama_stack.providers.utils.memory.vector_store import (
 
 from .config import ChromaVectorIOConfig as RemoteChromaVectorIOConfig
 
-log = logging.getLogger(__name__)
-
 ChromaClientType = chromadb.api.AsyncClientAPI | chromadb.api.ClientAPI
 
 VERSION = "v3"

llama_stack/providers/remote/vector_io/milvus/milvus.py

@@ -5,7 +5,6 @@
 # the root directory of this source tree.
 
 import asyncio
-import logging
 import os
 import re
 from typing import Any
@@ -21,6 +20,7 @@ from llama_stack.apis.vector_io import (
     QueryChunksResponse,
     VectorIO,
 )
+from llama_stack.log import get_logger
 from llama_stack.providers.datatypes import VectorDBsProtocolPrivate
 from llama_stack.providers.inline.vector_io.milvus import MilvusVectorIOConfig as InlineMilvusVectorIOConfig
 from llama_stack.providers.utils.kvstore import kvstore_impl
@@ -33,7 +33,7 @@ from llama_stack.providers.utils.memory.vector_store import (
 
 from .config import MilvusVectorIOConfig as RemoteMilvusVectorIOConfig
 
-logger = logging.getLogger(__name__)
+logger = get_logger(name=__name__, category="core")
 
 VERSION = "v3"
 VECTOR_DBS_PREFIX = f"vector_dbs:milvus:{VERSION}::"

llama_stack/providers/remote/vector_io/pgvector/pgvector.py

@@ -4,7 +4,6 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-import logging
 from typing import Any
 
 import psycopg2
@@ -21,6 +20,9 @@ from llama_stack.apis.vector_io import (
     QueryChunksResponse,
     VectorIO,
 )
+from llama_stack.log import get_logger
+
+log = get_logger(__name__, category="core")
 from llama_stack.providers.datatypes import Api, VectorDBsProtocolPrivate
 from llama_stack.providers.utils.kvstore import kvstore_impl
 from llama_stack.providers.utils.kvstore.api import KVStore
@@ -32,8 +34,6 @@ from llama_stack.providers.utils.memory.vector_store import (
 
 from .config import PGVectorVectorIOConfig
 
-log = logging.getLogger(__name__)
-
 VERSION = "v3"
 VECTOR_DBS_PREFIX = f"vector_dbs:pgvector:{VERSION}::"
 VECTOR_INDEX_PREFIX = f"vector_index:pgvector:{VERSION}::"

llama_stack/providers/remote/vector_io/qdrant/qdrant.py

@@ -4,7 +4,6 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-import logging
 import uuid
 from typing import Any
 
@@ -29,6 +28,9 @@ from llama_stack.apis.vector_io import (
     VectorStoreObject,
     VectorStoreSearchResponsePage,
 )
+from llama_stack.log import get_logger
+
+log = get_logger(__name__, category="core")
 from llama_stack.providers.datatypes import Api, VectorDBsProtocolPrivate
 from llama_stack.providers.inline.vector_io.qdrant import QdrantVectorIOConfig as InlineQdrantVectorIOConfig
 from llama_stack.providers.utils.memory.vector_store import (
@@ -38,7 +40,6 @@ from llama_stack.providers.utils.memory.vector_store import (
 
 from .config import QdrantVectorIOConfig as RemoteQdrantVectorIOConfig
 
-log = logging.getLogger(__name__)
 CHUNK_ID_KEY = "_chunk_id"
 
 

llama_stack/providers/remote/vector_io/weaviate/weaviate.py

@@ -4,7 +4,6 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 import json
-import logging
 from typing import Any
 
 import weaviate
@@ -18,6 +17,9 @@ from llama_stack.apis.files.files import Files
 from llama_stack.apis.vector_dbs import VectorDB
 from llama_stack.apis.vector_io import Chunk, QueryChunksResponse, VectorIO
 from llama_stack.distribution.request_headers import NeedsRequestProviderData
+from llama_stack.log import get_logger
+
+log = get_logger(__name__, category="core")
 from llama_stack.providers.datatypes import Api, VectorDBsProtocolPrivate
 from llama_stack.providers.utils.kvstore import kvstore_impl
 from llama_stack.providers.utils.kvstore.api import KVStore
@@ -28,8 +30,6 @@ from llama_stack.providers.utils.memory.vector_store import (
 
 from .config import WeaviateRequestProviderData, WeaviateVectorIOConfig
 
-log = logging.getLogger(__name__)
-
 VERSION = "v3"
 VECTOR_DBS_PREFIX = f"vector_dbs:weaviate:{VERSION}::"
 VECTOR_INDEX_PREFIX = f"vector_index:weaviate:{VERSION}::"

llama_stack/providers/utils/inference/embedding_mixin.py

@@ -5,7 +5,6 @@
 # the root directory of this source tree.
 
 import base64
-import logging
 import struct
 from typing import TYPE_CHECKING
 
@@ -27,7 +26,9 @@ from llama_stack.providers.utils.inference.prompt_adapter import interleaved_con
 EMBEDDING_MODELS = {}
 
 
-log = logging.getLogger(__name__)
+from llama_stack.log import get_logger
+
+logger = get_logger(name=__name__, category="inference")
 
 
 class SentenceTransformerEmbeddingMixin:
@@ -99,7 +100,7 @@ class SentenceTransformerEmbeddingMixin:
         if loaded_model is not None:
             return loaded_model
 
-        log.info(f"Loading sentence transformer for {model}...")
+        logger.info(f"Loading sentence transformer for {model}...")
         from sentence_transformers import SentenceTransformer
 
         loaded_model = SentenceTransformer(model)

llama_stack/providers/utils/inference/openai_compat.py

@@ -5,7 +5,6 @@
 # the root directory of this source tree.
 import base64
 import json
-import logging
 import struct
 import time
 import uuid
@@ -116,6 +115,7 @@ from llama_stack.apis.inference import (
 from llama_stack.apis.inference import (
     OpenAIChoice as OpenAIChatCompletionChoice,
 )
+from llama_stack.log import get_logger
 from llama_stack.models.llama.datatypes import (
     BuiltinTool,
     StopReason,
@@ -128,7 +128,7 @@ from llama_stack.providers.utils.inference.prompt_adapter import (
     decode_assistant_message,
 )
 
-logger = logging.getLogger(__name__)
+logger = get_logger(name=__name__, category="inference")
 
 
 class OpenAICompatCompletionChoiceDelta(BaseModel):

llama_stack/providers/utils/kvstore/mongodb/mongodb.py

@@ -4,17 +4,17 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-import logging
 from datetime import datetime
 
 from pymongo import AsyncMongoClient
 
+from llama_stack.log import get_logger
+
+logger = get_logger(name=__name__, category="core")
 from llama_stack.providers.utils.kvstore import KVStore
 
 from ..config import MongoDBKVStoreConfig
 
-log = logging.getLogger(__name__)
-
 
 class MongoDBKVStoreImpl(KVStore):
     def __init__(self, config: MongoDBKVStoreConfig):
@@ -34,7 +34,7 @@ class MongoDBKVStoreImpl(KVStore):
             self.conn = AsyncMongoClient(**conn_creds)
             self.collection = self.conn[self.config.db][self.config.collection_name]
         except Exception as e:
-            log.exception("Could not connect to MongoDB database server")
+            logger.exception("Could not connect to MongoDB database server")
             raise RuntimeError("Could not connect to MongoDB database server") from e
 
     def _namespaced_key(self, key: str) -> str:

llama_stack/providers/utils/kvstore/postgres/postgres.py

@@ -4,17 +4,18 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-import logging
 from datetime import datetime
 
 import psycopg2
 from psycopg2.extras import DictCursor
 
+from llama_stack.log import get_logger
+
+logger = get_logger(name=__name__, category="core")
+
 from ..api import KVStore
 from ..config import PostgresKVStoreConfig
 
-log = logging.getLogger(__name__)
-
 
 class PostgresKVStoreImpl(KVStore):
     def __init__(self, config: PostgresKVStoreConfig):
@@ -45,7 +46,7 @@ class PostgresKVStoreImpl(KVStore):
                 """
             )
         except Exception as e:
-            log.exception("Could not connect to PostgreSQL database server")
+            logger.exception("Could not connect to PostgreSQL database server")
             raise RuntimeError("Could not connect to PostgreSQL database server") from e
 
     def _namespaced_key(self, key: str) -> str:

llama_stack/providers/utils/memory/openai_vector_store_mixin.py

@@ -6,7 +6,6 @@
 
 import asyncio
 import json
-import logging
 import mimetypes
 import time
 import uuid
@@ -36,10 +35,11 @@ from llama_stack.apis.vector_io import (
     VectorStoreSearchResponse,
     VectorStoreSearchResponsePage,
 )
+from llama_stack.log import get_logger
 from llama_stack.providers.utils.kvstore.api import KVStore
 from llama_stack.providers.utils.memory.vector_store import content_from_data_and_mime_type, make_overlapped_chunks
 
-logger = logging.getLogger(__name__)
+logger = get_logger(name=__name__, category="core")
 
 # Constants for OpenAI vector stores
 CHUNK_MULTIPLIER = 5

llama_stack/providers/utils/memory/vector_store.py

@@ -5,7 +5,6 @@
 # the root directory of this source tree.
 import base64
 import io
-import logging
 import re
 import time
 from abc import ABC, abstractmethod
@@ -25,6 +24,9 @@ from llama_stack.apis.common.content_types import (
 from llama_stack.apis.tools import RAGDocument
 from llama_stack.apis.vector_dbs import VectorDB
 from llama_stack.apis.vector_io import Chunk, ChunkMetadata, QueryChunksResponse
+from llama_stack.log import get_logger
+
+logger = get_logger(name=__name__, category="memory")
 from llama_stack.models.llama.llama3.tokenizer import Tokenizer
 from llama_stack.providers.datatypes import Api
 from llama_stack.providers.utils.inference.prompt_adapter import (
@@ -32,8 +34,6 @@ from llama_stack.providers.utils.inference.prompt_adapter import (
 )
 from llama_stack.providers.utils.vector_io.chunk_utils import generate_chunk_id
 
-log = logging.getLogger(__name__)
-
 # Constants for reranker types
 RERANKER_TYPE_RRF = "rrf"
 RERANKER_TYPE_WEIGHTED = "weighted"
@@ -102,16 +102,16 @@ def content_from_data_and_mime_type(data: bytes | str, mime_type: str | None, en
             except UnicodeDecodeError as e:
                 if first_exception is None:
                     first_exception = e
-                log.warning(f"Decoding failed with {encoding}: {e}")
+                logger.warning(f"Decoding failed with {encoding}: {e}")
         # raise the origional exception, if we got here there was at least 1 exception
-        log.error(f"Could not decode data as any of {encodings_to_try}")
+        logger.error(f"Could not decode data as any of {encodings_to_try}")
         raise first_exception
 
     elif mime_type == "application/pdf":
         return parse_pdf(data)
 
     else:
-        log.error("Could not extract content from data_url properly.")
+        logger.error("Could not extract content from data_url properly.")
         return ""
 
 

llama_stack/providers/utils/telemetry/tracing.py

@@ -6,7 +6,7 @@
 
 import asyncio
 import contextvars
-import logging
+import logging  # allow-direct-logging
 import queue
 import random
 import threading

tests/integration/post_training/test_post_training.py

@@ -4,7 +4,6 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-import logging
 import sys
 import time
 import uuid
@@ -16,10 +15,9 @@ from llama_stack.apis.post_training import (
     LoraFinetuningConfig,
     TrainingConfig,
 )
+from llama_stack.log import get_logger
 
-# Configure logging
-logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s", force=True)
-logger = logging.getLogger(__name__)
+logger = get_logger(name=__name__, category="core")
 
 
 @pytest.fixture(autouse=True)

tests/integration/vector_io/test_openai_vector_stores.py

@@ -4,7 +4,6 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-import logging
 import time
 from io import BytesIO
 
@@ -15,8 +14,9 @@ from openai import OpenAI
 
 from llama_stack.apis.vector_io import Chunk
 from llama_stack.distribution.library_client import LlamaStackAsLibraryClient
+from llama_stack.log import get_logger
 
-logger = logging.getLogger(__name__)
+logger = get_logger(name=__name__, category="vector-io")
 
 
 def skip_if_provider_doesnt_support_openai_vector_stores(client_with_models):

tests/unit/providers/inference/test_remote_vllm.py

@@ -6,7 +6,7 @@
 
 import asyncio
 import json
-import logging
+import logging  # allow-direct-logging
 import threading
 import time
 from http.server import BaseHTTPRequestHandler, HTTPServer