Repository: SeldonIO/MLServer

Issue:
Support MLflow current protocol
As a follow-up to #167, it would be interesting to explore adding a custom endpoint to the `mlserver-mlflow` runtime which supports [MLflow's existing API](https://www.mlflow.org/docs/latest/models.html#deploy-mlflow-models). This would help reduce friction on user adoption of MLSever, as well as a temporary stopgap for users while they adopt the V2 protocol.


Discussion:

Resolving patch:
diff --git a/runtimes/mlflow/mlserver_mlflow/runtime.py b/runtimes/mlflow/mlserver_mlflow/runtime.py
--- a/runtimes/mlflow/mlserver_mlflow/runtime.py
+++ b/runtimes/mlflow/mlserver_mlflow/runtime.py
@@ -1,9 +1,29 @@
 import mlflow
 
+from io import StringIO
+from fastapi import Request, Response
+
+from mlflow.exceptions import MlflowException
+from mlflow.pyfunc.scoring_server import (
+    CONTENT_TYPES,
+    CONTENT_TYPE_CSV,
+    CONTENT_TYPE_JSON,
+    CONTENT_TYPE_JSON_SPLIT_ORIENTED,
+    CONTENT_TYPE_JSON_RECORDS_ORIENTED,
+    CONTENT_TYPE_JSON_SPLIT_NUMPY,
+    parse_csv_input,
+    infer_and_parse_json_input,
+    parse_json_input,
+    parse_split_oriented_json_input_to_numpy,
+    predictions_to_json,
+)
+
 from mlserver.types import InferenceRequest, InferenceResponse
 from mlserver.model import MLModel
 from mlserver.utils import get_model_uri
 from mlserver.codecs import get_decoded_or_raw
+from mlserver.handlers import custom_handler
+from mlserver.errors import InferenceError
 
 from .encoding import to_outputs
 
@@ -14,10 +34,68 @@ class MLflowRuntime(MLModel):
     models persisted with `joblib`.
     """
 
+    # TODO: Decouple from REST
+    @custom_handler(rest_path="/invocations")
+    async def invocations(self, request: Request) -> Response:
+        """
+        This custom handler is meant to mimic the behaviour of the existing
+        scoring server in MLflow.
+        For details about its implementation, please consult the original
+        implementation in the MLflow repository:
+
+            https://github.com/mlflow/mlflow/blob/master/mlflow/pyfunc/scoring_server/__init__.py
+        """
+        content_type = request.headers.get("content-type", None)
+        raw_data = await request.body()
+        as_str = raw_data.decode("utf-8")
+
+        if content_type == CONTENT_TYPE_CSV:
+            csv_input = StringIO(as_str)
+            data = parse_csv_input(csv_input=csv_input)
+        elif content_type == CONTENT_TYPE_JSON:
+            data = infer_and_parse_json_input(as_str, self._input_schema)
+        elif content_type == CONTENT_TYPE_JSON_SPLIT_ORIENTED:
+            data = parse_json_input(
+                json_input=StringIO(as_str),
+                orient="split",
+                schema=self._input_schema,
+            )
+        elif content_type == CONTENT_TYPE_JSON_RECORDS_ORIENTED:
+            data = parse_json_input(
+                json_input=StringIO(as_str),
+                orient="records",
+                schema=self._input_schema,
+            )
+        elif content_type == CONTENT_TYPE_JSON_SPLIT_NUMPY:
+            data = parse_split_oriented_json_input_to_numpy(as_str)
+        else:
+            content_type_error_message = (
+                "This predictor only supports the following content types, "
+                f"{CONTENT_TYPES}. Got '{content_type}'."
+            )
+            raise InferenceError(content_type_error_message)
+
+        try:
+            raw_predictions = self._model.predict(data)
+        except MlflowException as e:
+            raise InferenceError(e.message)
+        except Exception:
+            error_message = (
+                "Encountered an unexpected error while evaluating the model. Verify"
+                " that the serialized input Dataframe is compatible with the model for"
+                " inference."
+            )
+            raise InferenceError(error_message)
+
+        result = StringIO()
+        predictions_to_json(raw_predictions, result)
+        return Response(content=result.getvalue(), media_type="application/json")
+
     async def load(self) -> bool:
         # TODO: Log info message
         model_uri = await get_model_uri(self._settings)
         self._model = mlflow.pyfunc.load_model(model_uri)
+        self._input_schema = self._model.metadata.get_input_schema()
 
         self.ready = True
         return self.ready


Test patch:
diff --git a/runtimes/mlflow/tests/rest/__init__.py b/runtimes/mlflow/tests/rest/__init__.py
new file mode 100644
diff --git a/runtimes/mlflow/tests/rest/conftest.py b/runtimes/mlflow/tests/rest/conftest.py
new file mode 100644
--- /dev/null
+++ b/runtimes/mlflow/tests/rest/conftest.py
@@ -0,0 +1,71 @@
+import pytest
+
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+
+from mlserver.handlers import DataPlane, ModelRepositoryHandlers
+from mlserver.registry import MultiModelRegistry
+from mlserver.rest import RESTServer
+from mlserver.repository import ModelRepository
+from mlserver import Settings
+
+from mlserver_mlflow import MLflowRuntime
+
+
+@pytest.fixture
+def model_repository(model_uri: str) -> ModelRepository:
+    return ModelRepository(model_uri)
+
+
+@pytest.fixture
+async def model_registry(runtime: MLflowRuntime) -> MultiModelRegistry:
+    model_registry = MultiModelRegistry()
+    await model_registry.load(runtime)
+    return model_registry
+
+
+@pytest.fixture
+def settings() -> Settings:
+    return Settings()
+
+
+@pytest.fixture
+def data_plane(settings: Settings, model_registry: MultiModelRegistry) -> DataPlane:
+    return DataPlane(settings=settings, model_registry=model_registry)
+
+
+@pytest.fixture
+def model_repository_handlers(
+    model_repository: ModelRepository, model_registry: MultiModelRegistry
+) -> ModelRepositoryHandlers:
+    return ModelRepositoryHandlers(
+        repository=model_repository, model_registry=model_registry
+    )
+
+
+@pytest.fixture
+async def rest_server(
+    settings: Settings,
+    data_plane: DataPlane,
+    model_repository_handlers: ModelRepositoryHandlers,
+    runtime: MLflowRuntime,
+) -> RESTServer:
+    server = RESTServer(
+        settings,
+        data_plane=data_plane,
+        model_repository_handlers=model_repository_handlers,
+    )
+
+    await server.add_custom_handlers(runtime)
+
+    return server
+
+
+@pytest.fixture
+def rest_app(rest_server: RESTServer) -> FastAPI:
+    return rest_server._app
+
+
+@pytest.fixture
+def rest_client(rest_app: FastAPI) -> TestClient:
+    return TestClient(rest_app)
diff --git a/runtimes/mlflow/tests/rest/test_invocations.py b/runtimes/mlflow/tests/rest/test_invocations.py
new file mode 100644
--- /dev/null
+++ b/runtimes/mlflow/tests/rest/test_invocations.py
@@ -0,0 +1,31 @@
+import pytest
+
+from typing import Union
+
+
+@pytest.mark.parametrize("content_type", ["", None, "application/pdf"])
+def test_invocations_invalid_content_type(rest_client, content_type: str):
+    response = rest_client.post("/invocations", headers={"Content-Type": content_type})
+
+    assert response.status_code == 400
+
+
+@pytest.mark.parametrize(
+    "content_type, payload",
+    [
+        ("application/json", {"columns": ["a"], "data": [1, 2, 3]}),
+        ("application/json; format=pandas-records", [{"a": 1}, {"a": 2}, {"a": 3}]),
+        ("application/json", {"instances": [1, 2, 3]}),
+        ("application/json", {"inputs": [1, 2, 3]}),
+    ],
+)
+def test_invocations(rest_client, content_type: str, payload: Union[list, dict]):
+    response = rest_client.post(
+        "/invocations", headers={"Content-Type": content_type}, json=payload
+    )
+
+    assert response.status_code == 200
+
+    y_pred = response.json()
+    assert isinstance(y_pred, list)
+    assert len(y_pred) == 3