Error Handling

When a model's prediction call encounters an exception, the error response will be displayed in the navio UI.

The MLflow model will respond with something like:

{
  "error_code": "BAD_REQUEST",
  "message": "Encountered an unexpected error while evaluating the model. Verify that the serialized input Dataframe is compatible with the model for inference.",
  "stack_trace": "Traceback (most recent call last):\n  File \"/miniconda/envs/custom_env/lib/python3.8/site-packages/mlflow/pyfunc/scoring_server/__init__.py\", line 264, in ..."
}

While the stack trace always contains relevant information, the message is usually not very useful since the model input is always provided in the expected format. Errors are usually due to faults in the model's logic.

In addition, MLflow does not log errors by default.

One convenient way to add custom error handling is by using a decorator:

import logging
from functools import wraps
import traceback


import mlflow
...


def prediction_call(predict_fn: callable) -> callable:
    logger = logging.getLogger('gunicorn.error')

    @wraps(predict_fn)
    def wrapper(*args, **kwargs) -> dict:
        try:
            return predict_fn(*args, **kwargs)
        except Exception as exc:
            logger.exception('Prediction call failed')
            return {
                'error_code': exc.__class__.__name__,
                'message': str(exc),
                'stack_trace': traceback.format_exc()
            }

    return wrapper


class Predictor(mlflow.pyfunc.PythonModel):

    ...

    @prediction_call
    def predict(self, context: mlflow.pyfunc.PythonModelContext,
                model_input: pd.DataFrame) -> dict:

        # Will result in a response containing:
        # {
        #     "error_code": "AssertionError",
        #     "message": "Model must always receive at least two rows of data",
        #     "stack_trace": <traceback-output>
        # }
        assert model_input.shape[0] >= 2, \
            'Model must always receive at least two rows of data'

        return self._predict(model_input)

tip

The prediction_call decorator can be imported from pynavio.