Custom Error Handling Doesn't work for MLFlow Server

[carloszanella]

Custom Error Doesn't work for MLFlow Server

Describe the bug

When defining custom error handling according to documentation and using an MLFlowServer the error handler doesn't work.

To reproduce

Any versions after the one that introduced the custom error handling, including current release. To reproduce, I created the minimal example below. Requires Seldon-core and mlflow.

The gist is that seldon expects the model_error_handler attribute defined in the model, but when MLFlowServer class loads a model from mlflow (mlflow.pyfunc.load_model), it doesn't load the model you defined. It instead loads a wrapper of a wrapper of the model you defined.

So even if I define in my model the error handling I want, the code that seldon uses from mlflow is not my code and it doesn't have the error handling function.

import shutil
from pathlib import Path
from unittest.mock import Mock

import flask
import mlflow
from flask import jsonify
from mlflow.pyfunc import PythonModel
from seldon_core.metrics import SeldonMetrics
from seldon_core.wrapper import get_rest_microservice


class SeldonCustomError(Exception):
    """Defined according to docs"""

    status_code = 404

    def __init__(self, message, application_error_code, http_status_code):
        Exception.__init__(self)
        self.message = message
        if http_status_code is not None:
            self.status_code = http_status_code
        self.application_error_code = application_error_code

    def to_dict(self):
        res = {
            "status": {
                "status": self.status_code,
                "message": self.message,
                "app_code": self.application_error_code,
            }
        }
        return res


class MLFlowPythonModelExample(PythonModel):
    """A model that the predict method raises a custom error."""

    def __init__(self, error_code):
        self._test_error_code = error_code

    def predict(self, context, model_input):
        raise SeldonCustomError("Custom error", 1, self._test_error_code)

    model_error_handler = flask.Blueprint("error_handlers", __name__)

    @model_error_handler.app_errorhandler(SeldonCustomError)
    def handle_custom_error(error):
        response = jsonify(error.to_dict())
        response.status_code = error.status_code
        return response


def _register_model_to_mlflow(model: MLFlowPythonModelExample) -> str:
    """Test setup. Saves a dummy model to mlflow running in a local registry. Returns model path"""
    mlflow_dir = Path.cwd() / "mlflow_dir"
    if mlflow_dir.exists():
        shutil.rmtree(str(mlflow_dir))  # cleanup

    mlflow.set_tracking_uri(str(mlflow_dir))

    exp_id = mlflow.create_experiment("test-experiment")
    with mlflow.start_run(experiment_id=exp_id) as run:
        mlflow.pyfunc.log_model("model", python_model=model)

    return run.info.artifact_uri + "/model"


def test_mlflow_seldon_minimal_example():
    custom_error = 400
    test_model = MLFlowPythonModelExample(custom_error)
    model_path = _register_model_to_mlflow(test_model)
    loaded_mlflow_model = mlflow.pyfunc.load_model(model_path)

    metrics = Mock(spec=SeldonMetrics)
    app = get_rest_microservice(loaded_mlflow_model, metrics)
    client = app.test_client()

    response = client.get('/predict?json={"data":{"ndarray":[1,2]}}')

    assert response.status_code == custom_error
    assert response.json["status"]["app_code"] == 1

Expected behaviour

I actually dislike how mlflow is implemented here, but I think Seldon should support it nonetheless. I expect seldon to know how to work with mlflow in a way that we can define custom error handling in MLflow models too. I have a small, contained, suggestion for this fix. It is not pretty but I think other changes would require bigger refactoring. Let me know what you think.

Proposed solution

1. add self.model_error_handler as None in init
2. when loading, check if "model_error_handler" is present in the underlying model. If so, assign the attribute defined in previous step to it

class MLFlowServer(SeldonComponent):
    def __init__(self, model_uri: str, xtype: str = "ndarray"):
        super().__init__()
        ...
        self.model_error_handler = None

    def load(self):
        logger.info(f"Downloading model from {self.model_uri}")
        ...

        if hasattr(self._model._model_impl.python_model, "model_error_handler"):
            self.model_error_handler = (
                self._model._model_impl.python_model.model_error_handler
            )

[adriangonz]

Hey @carloszanella , thanks a lot for taking the time to raise this issue and documenting it properly.

We're currently on the process of moving towards MLServer, our next-gen inference server, where we've been working closely with the MLflow team to improve the existing integration with their artifacts. You can see an end-to-end example on how to use it "standalone" (i.e. locally) here, although you should also be able to enable it in Seldon Core by setting the protocol: kfserving field in your SeldonDeployment.

Given that all new developments will generally be added towards the new inference server, we'd be really keen to also support this use case there. @carloszanella would you be able to give it a go and see what's the behaviour there? I imagine that the only change required on your custom code would be raising a MLServerError exception, which should then be caught by MLServer and processed correctly.

Conscious that it's quite an information dump, so happy to answer any questions!