Repository: SeldonIO/MLServer

Issue:
Update to MLflow 2.0
MLflow 2.0 has updated their scoring protocol. Therefore, to ensure drop-in compatibility, we should update the MLflow runtime to follow the new updated spec. 


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
@@ -2,20 +2,15 @@
 
 from io import StringIO
 from fastapi import Depends, Header, Request, Response
-from typing import Optional
 
+from mlflow.version import VERSION
 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,
 )
 
@@ -60,12 +55,37 @@ async def ping(self) -> str:
         """
         return "\n"
 
+    @custom_handler(rest_path="/health", rest_method="GET")
+    async def health(self) -> str:
+        """
+        This custom handler is meant to mimic the behaviour of the existing
+        health endpoint in MLflow's local dev server.
+        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
+        """
+        return "\n"
+
+    # TODO: Decouple from REST
+    @custom_handler(rest_path="/version", rest_method="GET")
+    async def mlflow_version(self) -> str:
+        """
+        This custom handler is meant to mimic the behaviour of the existing
+        version endpoint in MLflow's local dev server.
+        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
+        """
+        return VERSION
+
     # TODO: Decouple from REST
     @custom_handler(rest_path="/invocations")
     async def invocations(
         self,
         raw_body: str = Depends(_get_raw_body),
-        content_type: Optional[str] = Header(default=None),
+        content_type: str = Header(default=""),
     ) -> Response:
         """
         This custom handler is meant to mimic the behaviour of the existing
@@ -75,31 +95,43 @@ async def invocations(
 
             https://github.com/mlflow/mlflow/blob/master/mlflow/pyfunc/scoring_server/__init__.py
         """
-        if content_type == CONTENT_TYPE_CSV:
+        # Content-Type can include other attributes like CHARSET
+        # Content-type RFC: https://datatracker.ietf.org/doc/html/rfc2045#section-5.1
+        # TODO: Suport ";" in quoted parameter values
+        type_parts = content_type.split(";")
+        type_parts = list(map(str.strip, type_parts))
+        mime_type = type_parts[0]
+        parameter_value_pairs = type_parts[1:]
+        parameter_values = {}
+        for parameter_value_pair in parameter_value_pairs:
+            (key, _, value) = parameter_value_pair.partition("=")
+            parameter_values[key] = value
+
+        charset = parameter_values.get("charset", "utf-8").lower()
+        if charset != "utf-8":
+            raise InferenceError("The scoring server only supports UTF-8")
+
+        unexpected_content_parameters = set(parameter_values.keys()).difference(
+            {"charset"}
+        )
+        if unexpected_content_parameters:
+            err_message = (
+                f"Unrecognized content type parameters: "
+                f"{', '.join(unexpected_content_parameters)}."
+            )
+            raise InferenceError(err_message)
+
+        if mime_type == CONTENT_TYPE_CSV:
             csv_input = StringIO(raw_body)
-            data = parse_csv_input(csv_input=csv_input)
-        elif content_type == CONTENT_TYPE_JSON:
+            data = parse_csv_input(csv_input=csv_input, schema=self._input_schema)
+        elif mime_type == CONTENT_TYPE_JSON:
             data = infer_and_parse_json_input(raw_body, self._input_schema)
-        elif content_type == CONTENT_TYPE_JSON_SPLIT_ORIENTED:
-            data = parse_json_input(
-                json_input=StringIO(raw_body),
-                orient="split",
-                schema=self._input_schema,
-            )
-        elif content_type == CONTENT_TYPE_JSON_RECORDS_ORIENTED:
-            data = parse_json_input(
-                json_input=StringIO(raw_body),
-                orient="records",
-                schema=self._input_schema,
-            )
-        elif content_type == CONTENT_TYPE_JSON_SPLIT_NUMPY:
-            data = parse_split_oriented_json_input_to_numpy(raw_body)
         else:
-            content_type_error_message = (
+            err_message = (
                 "This predictor only supports the following content types, "
                 f"{CONTENT_TYPES}. Got '{content_type}'."
             )
-            raise InferenceError(content_type_error_message)
+            raise InferenceError(err_message)
 
         try:
             raw_predictions = self._model.predict(data)


Test patch:
diff --git a/runtimes/mlflow/tests/conftest.py b/runtimes/mlflow/tests/conftest.py
--- a/runtimes/mlflow/tests/conftest.py
+++ b/runtimes/mlflow/tests/conftest.py
@@ -3,6 +3,7 @@
 import pytest
 import asyncio
 import numpy as np
+import pandas as pd
 
 from sklearn.dummy import DummyClassifier
 from pytorch_lightning import Trainer
@@ -33,7 +34,7 @@ def event_loop():
 @pytest.fixture
 def dataset() -> tuple:
     n = 4
-    X = np.random.rand(n)
+    X = pd.DataFrame({"foo": np.random.rand(n)})
     y = np.random.rand(n)
 
     return X, y
@@ -44,8 +45,6 @@ def model_signature(dataset: tuple) -> ModelSignature:
     X, y = dataset
     signature = infer_signature(X, y)
 
-    signature.inputs.inputs[0]._name = "foo"
-
     return signature
 
 
diff --git a/runtimes/mlflow/tests/rest/test_endpoints.py b/runtimes/mlflow/tests/rest/test_endpoints.py
new file mode 100644
--- /dev/null
+++ b/runtimes/mlflow/tests/rest/test_endpoints.py
@@ -0,0 +1,18 @@
+import pytest
+
+from mlflow.version import VERSION
+
+
+@pytest.mark.parametrize("endpoint", ["/ping", "/health"])
+async def test_ping(rest_client, endpoint: str):
+    response = await rest_client.get(endpoint)
+
+    assert response.status_code == 200
+    assert response.json() == "\n"
+
+
+async def test_version(rest_client):
+    response = await rest_client.get("/version")
+
+    assert response.status_code == 200
+    assert response.json() == VERSION
diff --git a/runtimes/mlflow/tests/rest/test_invocations.py b/runtimes/mlflow/tests/rest/test_invocations.py
--- a/runtimes/mlflow/tests/rest/test_invocations.py
+++ b/runtimes/mlflow/tests/rest/test_invocations.py
@@ -19,22 +19,53 @@ async def test_invocations_invalid_content_type(
 @pytest.mark.parametrize(
     "content_type, payload",
     [
-        ("application/json", {"columns": ["foo"], "data": [1, 2, 3]}),
-        (
-            "application/json; format=pandas-records",
-            [{"foo": 1}, {"foo": 2}, {"foo": 3}],
-        ),
         ("application/json", {"instances": [1, 2, 3]}),
         ("application/json", {"inputs": [1, 2, 3]}),
+        ("application/json", {"inputs": {"foo": [1, 2, 3]}}),
+    ],
+)
+async def test_invocations_tensor(
+    rest_client, content_type: str, payload: Union[list, dict]
+):
+    response = await rest_client.post(
+        "/invocations", headers={"Content-Type": content_type}, json=payload
+    )
+
+    assert response.status_code == 200
+
+    res = response.json()
+    assert "predictions" in res
+
+    y_pred = res["predictions"]
+    assert isinstance(y_pred, list)
+    assert len(y_pred) == 3
+
+
+@pytest.mark.parametrize(
+    "content_type, payload",
+    [
+        (
+            "application/json",
+            {"dataframe_split": {"columns": ["foo"], "data": [1, 2, 3]}},
+        ),
+        (
+            "application/json",
+            {"dataframe_records": [{"foo": 1}, {"foo": 2}, {"foo": 3}]},
+        ),
     ],
 )
-async def test_invocations(rest_client, content_type: str, payload: Union[list, dict]):
+async def test_invocations_dataframe(
+    rest_client, content_type: str, payload: Union[list, dict]
+):
     response = await rest_client.post(
         "/invocations", headers={"Content-Type": content_type}, json=payload
     )
 
     assert response.status_code == 200
 
-    y_pred = response.json()
+    res = response.json()
+    assert "predictions" in res
+
+    y_pred = res["predictions"]
     assert isinstance(y_pred, list)
-    assert len(y_pred) == 1
+    assert len(y_pred) == 3
diff --git a/runtimes/mlflow/tests/rest/test_ping.py b/runtimes/mlflow/tests/rest/test_ping.py
deleted file mode 100644
--- a/runtimes/mlflow/tests/rest/test_ping.py
+++ /dev/null
@@ -1,5 +0,0 @@
-async def test_invocations(rest_client):
-    response = await rest_client.get("/ping")
-
-    assert response.status_code == 200
-    assert response.json() == "\n"
diff --git a/runtimes/mlflow/tests/test_runtime.py b/runtimes/mlflow/tests/test_runtime.py
--- a/runtimes/mlflow/tests/test_runtime.py
+++ b/runtimes/mlflow/tests/test_runtime.py
@@ -4,7 +4,7 @@
 
 from typing import Any
 
-from mlserver.codecs import NumpyCodec
+from mlserver.codecs import NumpyCodec, PandasCodec
 from mlserver.types import (
     InferenceRequest,
     Parameters,
@@ -16,7 +16,6 @@
 from mlflow.models.signature import ModelSignature
 
 from mlserver_mlflow import MLflowRuntime
-from mlserver_mlflow.codecs import TensorDictCodec
 
 
 def test_load(runtime: MLflowRuntime):
@@ -159,4 +158,4 @@ async def test_metadata(runtime: MLflowRuntime, model_signature: ModelSignature)
     assert len(model_signature.outputs.inputs) == len(metadata.outputs)
 
     assert metadata.parameters is not None
-    assert metadata.parameters.content_type == TensorDictCodec.ContentType
+    assert metadata.parameters.content_type == PandasCodec.ContentType