feat(logging): Distinguish logs from different models

docs/examples/model-repository/README.ipynb

@@ -75,7 +75,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "As we can, the repository lists 2 models (i.e. `mushroom-xgboost` and `mnist-svm`).\n",
+    "As we can see, the repository lists 2 models (i.e. `mushroom-xgboost` and `mnist-svm`).\n",
     "Note that the state for both is set to `READY`.\n",
     "This means that both models are loaded, and thus ready for inference."
    ]

mlserver/context.py

@@ -0,0 +1,23 @@
+from contextlib import contextmanager
+from contextvars import ContextVar
+
+from .settings import ModelSettings
+
+model_name_var: ContextVar[str] = ContextVar("model_name")
+model_version_var: ContextVar[str] = ContextVar("model_version")
+
+
+@contextmanager
+def model_context(model_settings: ModelSettings):
+    model_name_token = model_name_var.set(model_settings.name)
+
+    model_version = ""
+    if model_settings.version:
+        model_version = model_settings.version
+    model_version_token = model_version_var.set(model_version)
+
+    try:
+        yield
+    finally:
+        model_name_var.reset(model_name_token)
+        model_version_var.reset(model_version_token)

mlserver/handlers/dataplane.py

@@ -5,7 +5,7 @@
 from typing import Optional
 
 from ..errors import ModelNotReady
-from ..metrics import model_context
+from ..context import model_context
 from ..settings import Settings
 from ..registry import MultiModelRegistry
 from ..types import (
@@ -73,7 +73,8 @@ async def model_metadata(
     ) -> MetadataModelResponse:
         model = await self._model_registry.get_model(name, version)
         # TODO: Make await optional for sync methods
-        return await model.metadata()
+        with model_context(model.settings):
+            return await model.metadata()
 
     async def infer(
         self,

mlserver/logging.py

@@ -2,15 +2,16 @@
 import logging
 import sys
 
-from logging import Formatter, StreamHandler
+from logging import StreamHandler
 from pathlib import Path
 from typing import Optional, Dict, Union
 import logging.config
 
+from .context import model_name_var, model_version_var
 from .settings import Settings
 
 LoggerName = "mlserver"
-LoggerFormat = "%(asctime)s [%(name)s] %(levelname)s - %(message)s"
+LoggerFormat = "%(asctime)s [%(name)s]%(model)s %(levelname)s - %(message)s"
 
 logger = logging.getLogger(LoggerName)
 
@@ -35,13 +36,35 @@ def apply_logging_file(logging_settings: Union[str, Dict]):
         logger.warning("Unable to parse logging_settings.")
 
 
+class ModelLoggerFormatter(logging.Formatter):
+    """
+    A logging formatter that uses context variables to inject
+    the model name and version in the log message.
+    """
+
+    @staticmethod
+    def _fmt_model(name: str, version: str) -> str:
+        if not name:
+            return ""
+        model_fmt = f"{name}:{version}" if version else name
+        return f"[{model_fmt}]"
+
+    def format(self, record: logging.LogRecord) -> str:
+        model_name = model_name_var.get("")
+        model_version = model_version_var.get("")
+
+        record.model = self.__class__._fmt_model(model_name, model_version)
+
+        return super().format(record)
+
+
 def configure_logger(settings: Optional[Settings] = None):
     logger = get_logger()
 
     # Don't add handler twice
     if not logger.handlers:
         stream_handler = StreamHandler(sys.stdout)
-        formatter = Formatter(LoggerFormat)
+        formatter = ModelLoggerFormatter(LoggerFormat)
         stream_handler.setFormatter(formatter)
 
         logger.addHandler(stream_handler)

mlserver/metrics/__init__.py

@@ -1,12 +1,11 @@
 from .server import MetricsServer
 from .prometheus import configure_metrics
-from .context import model_context, register, log
+from .context import register, log
 from .registry import REGISTRY
 
 __all__ = [
     "MetricsServer",
     "configure_metrics",
-    "model_context",
     "register",
     "log",
     "REGISTRY",

mlserver/metrics/context.py

@@ -1,34 +1,13 @@
-from contextvars import ContextVar
-from contextlib import contextmanager
 from prometheus_client import Histogram
 
-from ..settings import ModelSettings
 from .registry import REGISTRY
 from .errors import InvalidModelContext
-
-model_name_var: ContextVar[str] = ContextVar("model_name")
-model_version_var: ContextVar[str] = ContextVar("model_version")
+from ..context import model_name_var, model_version_var
 
 SELDON_MODEL_NAME_LABEL = "model_name"
 SELDON_MODEL_VERSION_LABEL = "model_version"
 
 
-@contextmanager
-def model_context(model_settings: ModelSettings):
-    model_name_token = model_name_var.set(model_settings.name)
-
-    model_version = ""
-    if model_settings.version:
-        model_version = model_settings.version
-    model_version_token = model_version_var.set(model_version)
-
-    try:
-        yield
-    finally:
-        model_name_var.reset(model_name_token)
-        model_version_var.reset(model_version_token)
-
-
 def register(name: str, description: str) -> Histogram:
     """
     Registers a new metric with its description.

mlserver/parallel/worker.py

@@ -12,7 +12,8 @@
 from ..utils import install_uvloop_event_loop, schedule_with_callback
 from ..logging import configure_logger
 from ..settings import Settings
-from ..metrics import configure_metrics, model_context
+from ..metrics import configure_metrics
+from ..context import model_context
 from ..env import Environment
 
 from .messages import (

mlserver/registry.py

@@ -4,6 +4,7 @@
 from itertools import chain
 from functools import cmp_to_key
 
+from .context import model_context
 from .model import MLModel
 from .errors import ModelNotFound
 from .logging import logger
@@ -142,10 +143,11 @@ async def load(self, model_settings: ModelSettings) -> MLModel:
 
         new_model = self._model_initialiser(model_settings)
 
-        if previous_loaded_model:
-            await self._reload_model(previous_loaded_model, new_model)
-        else:
-            await self._load_model(new_model)
+        with model_context(model_settings):
+            if previous_loaded_model:
+                await self._reload_model(previous_loaded_model, new_model)
+            else:
+                await self._load_model(new_model)
 
         return new_model
 
@@ -202,27 +204,28 @@ async def unload_version(self, version: Optional[str] = None):
         model = await self.get_model(version)
         await self._unload_model(model)
 
-        model_msg = f"model '{model.name}'"
-        if version:
-            model_msg = f"version {version} of {model_msg}"
+    async def _unload_model(self, model: MLModel):
+        with model_context(model.settings):
+            # NOTE: Every callback needs to run to ensure one doesn't block the
+            # others
+            await asyncio.gather(
+                *[callback(model) for callback in self._on_model_unload],
+                return_exceptions=True,
+            )
 
-        logger.info(f"Unloaded {model_msg} succesfully.")
+            if model.version:
+                del self._versions[model.version]
 
-    async def _unload_model(self, model: MLModel):
-        # NOTE: Every callback needs to run to ensure one doesn't block the
-        # others
-        await asyncio.gather(
-            *[callback(model) for callback in self._on_model_unload],
-            return_exceptions=True,
-        )
+            if model == self.default:
+                self._clear_default()
 
-        if model.version:
-            del self._versions[model.version]
+            model.ready = not await model.unload()
 
-        if model == self.default:
-            self._clear_default()
+            model_msg = f"model '{model.name}'"
+            if model.version:
+                model_msg = f"version {model.version} of {model_msg}"
 
-        model.ready = not await model.unload()
+            logger.info(f"Unloaded {model_msg} succesfully.")
 
     def _find_model(self, version: Optional[str] = None) -> Optional[MLModel]:
         if version:

runtimes/alibi-detect/tests/conftest.py

@@ -10,7 +10,7 @@
 from alibi_detect.od import OutlierVAE
 from alibi_detect.saving import save_detector
 
-from mlserver.metrics.context import model_context
+from mlserver.context import model_context
 from mlserver.settings import ModelSettings, ModelParameters
 from mlserver.types import InferenceRequest
 from mlserver.utils import install_uvloop_event_loop

tests/batching/conftest.py

@@ -7,7 +7,7 @@
 from mlserver.utils import generate_uuid
 from mlserver.types import InferenceRequest, TensorData
 from mlserver.model import MLModel
-from mlserver.metrics.context import model_context
+from mlserver.context import model_context
 from mlserver.batching.adaptive import AdaptiveBatcher
 from mlserver.settings import ModelSettings
 

tests/metrics/test_context.py

@@ -13,9 +13,9 @@
     SELDON_MODEL_VERSION_LABEL,
     register,
     log,
-    model_context,
     _get_labels_from_context,
 )
+from mlserver.context import model_context
 
 from ..fixtures import SumModel
 

tests/test_context.py

@@ -0,0 +1,47 @@
+import pytest
+
+from typing import Optional
+
+from mlserver.settings import ModelSettings, ModelParameters
+from mlserver.context import model_context, model_name_var, model_version_var
+from .fixtures import SumModel
+
+
+@pytest.mark.parametrize(
+    "name, version, expected_version",
+    [
+        (
+            "foo",
+            "v1.0",
+            "v1.0",
+        ),
+        (
+            "foo",
+            "",
+            "",
+        ),
+        ("foo", None, ""),
+    ],
+)
+def test_model_context(name: str, version: Optional[str], expected_version: str):
+    model_settings = ModelSettings(
+        name=name, implementation=SumModel, parameters=ModelParameters(version=version)
+    )
+
+    with pytest.raises(LookupError):
+        _ = model_name_var.get()
+
+    with pytest.raises(LookupError):
+        _ = model_version_var.get()
+
+    with model_context(model_settings):
+        var_name = model_name_var.get()
+        var_version = model_version_var.get()
+        assert var_name == name
+        assert var_version == expected_version
+
+    with pytest.raises(LookupError):
+        _ = model_name_var.get()
+
+    with pytest.raises(LookupError):
+        _ = model_version_var.get()

tests/test_logging.py

@@ -1,7 +1,57 @@
 import pytest
 
-from mlserver.logging import logger, configure_logger
-from mlserver.settings import Settings
+from mlserver import ModelSettings
+from mlserver.context import model_context
+from mlserver.logging import (
+    ModelLoggerFormatter,
+    LoggerFormat,
+    configure_logger,
+    logger,
+)
+from mlserver.settings import ModelParameters, Settings
+from tests.fixtures import SumModel
+from logging import INFO
+
+
+@pytest.mark.parametrize(
+    "name, version, expected_fmt",
+    [
+        (
+            "foo",
+            "v1.0",
+            "[foo:v1.0]",
+        ),
+        (
+            "foo",
+            "",
+            "[foo]",
+        ),
+        (
+            "foo",
+            None,
+            "[foo]",
+        ),
+    ],
+)
+def test_model_logging_formatter(name, version, expected_fmt, caplog):
+    caplog.handler.setFormatter(ModelLoggerFormatter(LoggerFormat))
+    caplog.set_level(INFO)
+
+    model_settings = ModelSettings(
+        name=name, implementation=SumModel, parameters=ModelParameters(version=version)
+    )
+
+    logger.info("Before model context")
+    with model_context(model_settings):
+        logger.info("Inside model context")
+    logger.info("After model context")
+
+    log_records = caplog.text.strip().split("\n")
+    assert len(caplog.messages) == 3
+
+    assert expected_fmt not in log_records[0]
+    assert expected_fmt in log_records[1]
+    assert expected_fmt not in log_records[2]
 
 
 @pytest.mark.parametrize("debug", [True, False])