Fixed logging when running the deployed model; fixed model params signature not present

Author: Lucashsmello

datamint/mlflow/__init__.py

@@ -15,6 +15,7 @@
 if mlflow_utils.is_tracking_uri_set():
     _LOGGER.warning("MLflow tracking URI is already set before patching get_tracking_uri.")
 
+
 @wraps(_original_get_tracking_uri)
 def _patched_get_tracking_uri(*args, **kwargs):
     """Patched version of get_tracking_uri that ensures MLflow environment is set up first.
@@ -47,6 +48,56 @@ def _patched_get_tracking_uri(*args, **kwargs):
 # Replace the original function with our patched version
 mlflow_utils.get_tracking_uri = _patched_get_tracking_uri
 
+_ALREADY_CONFIGURED_LOGGING = False
+
+
+def _configure_mlflow_loggers():
+    global _ALREADY_CONFIGURED_LOGGING
+    if _ALREADY_CONFIGURED_LOGGING:
+        return
+
+    from mlflow.environment_variables import MLFLOW_LOGGING_LEVEL
+    from mlflow.utils.logging_utils import SuppressLogFilter, get_mlflow_log_level
+    import logging.config
+    import rich.logging
+    import os
+
+    if 'MLFLOW_SCORING_SERVER_REQUEST_TIMEOUT' not in os.environ:
+        # probably not running in mlflow server, so no need to configure our mlflow loggers
+        return
+
+    _ALREADY_CONFIGURED_LOGGING = True
+
+    logging.config.dictConfig(
+        {
+            "version": 1,
+            "disable_existing_loggers": False,
+            "handlers": {
+                "datamint_mlflow_handler": {
+                    "class": "rich.logging.RichHandler",
+                    "filters": ["suppress_in_thread"],
+                },
+            },
+            "loggers": {
+                'datamint.mlflow': {
+                    "handlers": ["datamint_mlflow_handler"],
+                    "level": get_mlflow_log_level(),
+                    "propagate": False,
+                },
+            },
+            "filters": {
+                "suppress_in_thread": {
+                    "()": SuppressLogFilter,
+                }
+            },
+        }
+    )
+    _LOGGER.info("Configured MLflow loggers with RichHandler and level %s", get_mlflow_log_level())
+
+try:
+    _configure_mlflow_loggers()
+except Exception as e:
+    _LOGGER.error("Failed to configure MLflow loggers: %s", e)
 
 if TYPE_CHECKING:
     from .flavors.model import DatamintModel
@@ -62,4 +113,5 @@ def _patched_get_tracking_uri(*args, **kwargs):
         },
     )
 
-__all__ = ['set_project', 'setup_mlflow_environment', 'ensure_mlflow_configured', 'DatamintModel']
+
+__all__ = ['set_project', 'setup_mlflow_environment', 'ensure_mlflow_configured', 'DatamintModel']

datamint/mlflow/flavors/datamint_flavor.py

@@ -11,8 +11,23 @@
 
 FLAVOR_NAME = 'datamint'
 
-_LOGGER = logging.getLogger(__name__)
 
+def _process_signature(signature: ModelSignature | None) -> ModelSignature:
+    from mlflow.types import ParamSchema, ParamSpec
+
+    # Define inference parameters
+    params_schema = ParamSchema(
+        [
+            ParamSpec("mode", "string", "default"),  # Default mode
+        ]
+    )
+
+    if signature is None:
+        signature = ModelSignature(params=params_schema)
+    else:
+        signature.params = params_schema
+    return signature
+    
 
 def save_model(datamint_model: DatamintModel,
                path,
@@ -77,6 +92,8 @@ def _get_req_name(req):
 
     datamint_model._clear_linked_models_cache()
 
+    signature = _process_signature(signature)
+
     return mlflow.pyfunc.save_model(
         path=path,
         python_model=datamint_model,

datamint/mlflow/flavors/model.py

@@ -399,6 +399,7 @@ def predict(self,
 
         # Parse and validate mode
         mode = self._parse_mode(model_input=model_input, params=params)
+        logger.info(f"Received prediction request with {len(model_input)} resources and params {params} with mode '{mode.value}'")
 
         # Route to appropriate prediction method
         try:
@@ -408,7 +409,7 @@ def predict(self,
                     mode = PredictionMode.DEFAULT
                 else:
                     raise NotImplementedError
-            logger.debug(f"Routing to '{mode.value}' mode for {len(model_input)} resources")
+            logger.info(f"Routing to '{mode.value}' mode for {len(model_input)} resources")
             result = self._route_prediction(model_input, mode, params)
 
             # Apply common post-processing

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "datamint"
 description = "A library for interacting with the Datamint API, designed for efficient data management, processing and Deep Learning workflows."
-version = "2.10.1"
+version = "2.10.2"
 dynamic = ["dependencies"]
 requires-python = ">=3.10"
 readme = "README.md"