Repository: SeldonIO/MLServer

Issue:
Support drift batch accumulation in Detect runtime
Drift models generally need to operate on a batch of inputs. This batch would usually get accumulated on-the-fly with a "moving window" over past requests that is kept within the runtime. 

Note that this type of batch is different from the adaptive batching support built into MLServer. On the latter, all request connections are held while the batch is getting formed (and only returned after batch inference occurs), whereas in the former requests are answered as they come (with the batch kept up-to-date in the background). 


Discussion:

Resolving patch:
diff --git a/mlserver/batching/__init__.py b/mlserver/batching/__init__.py
--- a/mlserver/batching/__init__.py
+++ b/mlserver/batching/__init__.py
@@ -1,4 +1,5 @@
+from .requests import BatchedRequests
 from .adaptive import AdaptiveBatcher
 from .hooks import load_batching
 
-__all__ = ["AdaptiveBatcher", "load_batching"]
+__all__ = ["AdaptiveBatcher", "load_batching", "BatchedRequests"]
diff --git a/runtimes/alibi-detect/mlserver_alibi_detect/runtime.py b/runtimes/alibi-detect/mlserver_alibi_detect/runtime.py
--- a/runtimes/alibi-detect/mlserver_alibi_detect/runtime.py
+++ b/runtimes/alibi-detect/mlserver_alibi_detect/runtime.py
@@ -1,14 +1,19 @@
-from mlserver import types
+import numpy as np
+
+from pydantic.error_wrappers import ValidationError
+from typing import Optional, List
+from pydantic import BaseSettings
+from functools import cached_property
+
+from alibi_detect.saving import load_detector
+
+from mlserver.batching import BatchedRequests
+from mlserver.types import InferenceRequest, InferenceResponse
 from mlserver.settings import ModelSettings
 from mlserver.model import MLModel
 from mlserver.codecs import NumpyCodec, NumpyRequestCodec
 from mlserver.utils import get_model_uri
-from alibi_detect.utils.saving import load_detector
 from mlserver.errors import MLServerError, InferenceError
-from pydantic.error_wrappers import ValidationError
-from typing import Optional
-from pydantic import BaseSettings
-import numpy as np
 
 ENV_PREFIX_ALIBI_DETECT_SETTINGS = "MLSERVER_MODEL_ALIBI_DETECT_"
 
@@ -23,7 +28,18 @@ class AlibiDetectSettings(BaseSettings):
     class Config:
         env_prefix = ENV_PREFIX_ALIBI_DETECT_SETTINGS
 
-    predict_parameters: Optional[dict] = {}
+    predict_parameters: dict = {}
+    """
+    Keyword parameters to pass to the detector's `predict` method.
+    """
+
+    batch_size: Optional[int] = None
+    """
+    Run the detector after accumulating a batch of size `batch_size`.
+    Note that this is different to MLServer's adaptive batching, since the rest
+    of requests will just return empty (i.e. instead of being hold until
+    inference runs for all of them).
+    """
 
 
 class AlibiDetectRuntime(MLModel):
@@ -33,10 +49,12 @@ class AlibiDetectRuntime(MLModel):
 
     def __init__(self, settings: ModelSettings):
         if settings.parameters is None:
-            self.alibi_detect_settings = AlibiDetectSettings()
+            self._ad_settings = AlibiDetectSettings()
         else:
             extra = settings.parameters.extra
-            self.alibi_detect_settings = AlibiDetectSettings(**extra)  # type: ignore
+            self._ad_settings = AlibiDetectSettings(**extra)  # type: ignore
+
+        self._batch: List[InferenceRequest] = []
         super().__init__(settings)
 
     async def load(self) -> bool:
@@ -57,41 +75,71 @@ async def load(self) -> bool:
         self.ready = True
         return self.ready
 
-    async def predict(self, payload: types.InferenceRequest) -> types.InferenceResponse:
+    async def predict(self, payload: InferenceRequest) -> InferenceResponse:
+        # If batch is not configured, run the detector and return the output
+        if not self._ad_settings.batch_size:
+            return self._detect(payload)
+
+        if len(self._batch) < self._ad_settings.batch_size:
+            self._batch.append(payload)
+
+        if len(self._batch) >= self._ad_settings.batch_size:
+            batched = self._get_batched_request()
+            return self._detect(batched)
+
+        return InferenceResponse(
+            model_name=self.name, model_version=self.version, outputs=[]
+        )
+
+    def _get_batched_request(self) -> InferenceRequest:
+        # Build requests dictionary with mocked IDs.
+        # This is required by the BatchedRequests util.
+        inference_requests = {}
+        for idx, inference_request in enumerate(self._batch):
+            inference_requests[str(idx)] = inference_request
+
+        batched = BatchedRequests(inference_requests)
+        self._batch = []
+        return batched.merged_request
+
+    def _detect(self, payload: InferenceRequest) -> InferenceResponse:
         input_data = self.decode_request(payload, default_codec=NumpyRequestCodec)
-        predict_kwargs = {}
-        if self.alibi_detect_settings.predict_parameters is not None:
-            predict_kwargs = self.alibi_detect_settings.predict_parameters
+        predict_kwargs = self._ad_settings.predict_parameters
 
         try:
             y = self._model.predict(np.array(input_data), **predict_kwargs)
+            return self._encode_response(y)
         except (ValueError, IndexError) as e:
             raise InferenceError(
                 f"Invalid predict parameters for model {self._settings.name}: {e}"
             ) from e
 
+    def _encode_response(self, y: dict) -> InferenceResponse:
         outputs = []
         for key in y["data"]:
             outputs.append(
                 NumpyCodec.encode_output(name=key, payload=np.array([y["data"][key]]))
             )
+
         # Add headers
         y["meta"]["headers"] = {
-            "x-seldon-alibi-type": self.get_alibi_type(),
-            "x-seldon-alibi-method": self.get_alibi_method(),
+            "x-seldon-alibi-type": self.alibi_type,
+            "x-seldon-alibi-method": self.alibi_method,
         }
-        return types.InferenceResponse(
+        return InferenceResponse(
             model_name=self.name,
             model_version=self.version,
             parameters=y["meta"],
             outputs=outputs,
         )
 
-    def get_alibi_method(self) -> str:
+    @cached_property
+    def alibi_method(self) -> str:
         module: str = type(self._model).__module__
         return module.split(".")[-1]
 
-    def get_alibi_type(self) -> str:
+    @cached_property
+    def alibi_type(self) -> str:
         module: str = type(self._model).__module__
         method = module.split(".")[-2]
         if method in ALIBI_MODULE_NAMES:


Test patch:
diff --git a/runtimes/alibi-detect/tests/conftest.py b/runtimes/alibi-detect/tests/conftest.py
--- a/runtimes/alibi-detect/tests/conftest.py
+++ b/runtimes/alibi-detect/tests/conftest.py
@@ -1,9 +1,23 @@
 import pytest
 import os
 import asyncio
+import numpy as np
+import tensorflow as tf
+
+from tensorflow.keras.layers import Dense, InputLayer
+from alibi_detect.cd import TabularDrift
+from alibi_detect.od import OutlierVAE
+from alibi_detect.saving import save_detector
+
+from mlserver.settings import ModelSettings, ModelParameters
 from mlserver.types import InferenceRequest
 from mlserver.utils import install_uvloop_event_loop
 
+from mlserver_alibi_detect import AlibiDetectRuntime
+
+tf.keras.backend.clear_session()
+
+P_VAL_THRESHOLD = 0.05
 
 TESTS_PATH = os.path.dirname(__file__)
 TESTDATA_PATH = os.path.join(TESTS_PATH, "testdata")
@@ -22,3 +36,110 @@ def event_loop():
 def inference_request() -> InferenceRequest:
     payload_path = os.path.join(TESTDATA_PATH, "inference-request.json")
     return InferenceRequest.parse_file(payload_path)
+
+
+@pytest.fixture
+def outlier_detector_settings(
+    outlier_detector_uri: str,
+) -> ModelSettings:
+    return ModelSettings(
+        name="alibi-detect-model",
+        implementation=AlibiDetectRuntime,
+        parameters=ModelParameters(
+            uri=outlier_detector_uri,
+            version="v1.2.3",
+            extra={
+                "predict_parameters": {
+                    "outlier_type": "instance",
+                    "return_feature_score": False,
+                    "return_instance_score": True,
+                }
+            },
+        ),
+    )
+
+
+@pytest.fixture
+def outlier_detector_uri(tmp_path: str) -> str:
+    X_ref = np.array([[1, 2, 3]])
+    n_features = X_ref.shape[1]
+    latent_dim = 2
+    encoder_net = tf.keras.Sequential(
+        [
+            InputLayer(input_shape=(n_features,)),
+            Dense(25, activation=tf.nn.relu),
+            Dense(10, activation=tf.nn.relu),
+            Dense(5, activation=tf.nn.relu),
+        ]
+    )
+
+    decoder_net = tf.keras.Sequential(
+        [
+            InputLayer(input_shape=(latent_dim,)),
+            Dense(5, activation=tf.nn.relu),
+            Dense(10, activation=tf.nn.relu),
+            Dense(25, activation=tf.nn.relu),
+            Dense(n_features, activation=None),
+        ]
+    )
+
+    od = OutlierVAE(
+        threshold=0.05,
+        score_type="mse",
+        encoder_net=encoder_net,
+        decoder_net=decoder_net,
+        latent_dim=latent_dim,
+        samples=5,
+    )
+
+    od.fit(X_ref, loss_fn=tf.keras.losses.mse, epochs=5, verbose=True)
+
+    detector_uri = os.path.join(tmp_path, "alibi-detector-artifacts")
+    save_detector(od, detector_uri)
+
+    return detector_uri
+
+
+@pytest.fixture
+async def outlier_detector(
+    outlier_detector_settings: ModelSettings,
+) -> AlibiDetectRuntime:
+    model = AlibiDetectRuntime(outlier_detector_settings)
+    await model.load()
+
+    return model
+
+
+@pytest.fixture
+def drift_detector_settings(
+    drift_detector_uri: str,
+) -> ModelSettings:
+    return ModelSettings(
+        name="alibi-detect-model",
+        implementation=AlibiDetectRuntime,
+        parameters=ModelParameters(
+            uri=drift_detector_uri,
+            version="v1.2.3",
+            extra={"predict_parameters": {"drift_type": "feature"}, "batch_size": 5},
+        ),
+    )
+
+
+@pytest.fixture
+def drift_detector_uri(tmp_path: str) -> str:
+    X_ref = np.array([[1, 2, 3]])
+
+    cd = TabularDrift(X_ref, p_val=P_VAL_THRESHOLD)
+
+    detector_uri = os.path.join(tmp_path, "alibi-detector-artifacts")
+    save_detector(cd, detector_uri)
+
+    return detector_uri
+
+
+@pytest.fixture
+async def drift_detector(drift_detector_settings: ModelSettings) -> AlibiDetectRuntime:
+    model = AlibiDetectRuntime(drift_detector_settings)
+    await model.load()
+
+    return model
diff --git a/runtimes/alibi-detect/tests/test_drift_detector.py b/runtimes/alibi-detect/tests/test_drift_detector.py
--- a/runtimes/alibi-detect/tests/test_drift_detector.py
+++ b/runtimes/alibi-detect/tests/test_drift_detector.py
@@ -1,74 +1,34 @@
-import pytest
-from mlserver.settings import ModelSettings, ModelParameters
-from mlserver.types import RequestInput, InferenceRequest
-from mlserver.codecs import CodecError
-from mlserver_alibi_detect import AlibiDetectRuntime
-from alibi_detect.cd import TabularDrift
-from alibi_detect.utils.saving import save_detector
-import os
 import numpy as np
 
-P_VAL_THRESHOLD = 0.05
-
-
-@pytest.fixture
-def alibi_detect_tabular_drift_model_settings(
-    alibi_detect_tabular_drift_model_uri: str,
-) -> ModelSettings:
-    return ModelSettings(
-        name="alibi-detect-model",
-        implementation=AlibiDetectRuntime,
-        parameters=ModelParameters(
-            uri=alibi_detect_tabular_drift_model_uri,
-            version="v1.2.3",
-            extra={"predict_parameters": {"drift_type": "feature"}},
-        ),
-    )
-
-
-@pytest.fixture
-def alibi_detect_tabular_drift_model_uri(tmp_path) -> str:
-    X_ref = np.array([[1, 2, 3]])
-
-    cd = TabularDrift(X_ref, p_val=P_VAL_THRESHOLD)
-
-    detector_uri = os.path.join(tmp_path, "alibi-detector-artifacts")
-    save_detector(cd, detector_uri)
-
-    return detector_uri
+from alibi_detect.cd import TabularDrift
 
+from mlserver.types import InferenceRequest
+from mlserver.codecs import NumpyRequestCodec
 
-@pytest.fixture
-async def alibi_detect_tabular_drift_model(
-    alibi_detect_tabular_drift_model_settings: ModelSettings,
-) -> AlibiDetectRuntime:
-    model = AlibiDetectRuntime(alibi_detect_tabular_drift_model_settings)
-    await model.load()
+from mlserver_alibi_detect import AlibiDetectRuntime
 
-    return model
+from .conftest import P_VAL_THRESHOLD
 
 
 async def test_load_folder(
-    alibi_detect_tabular_drift_model_uri: str,
-    alibi_detect_tabular_drift_model_settings: ModelSettings,
+    drift_detector: AlibiDetectRuntime,
 ):
-    alibi_detect_tabular_drift_model_settings.parameters.uri = (  # type: ignore
-        alibi_detect_tabular_drift_model_uri
-    )
-
-    model = AlibiDetectRuntime(alibi_detect_tabular_drift_model_settings)
-    await model.load()
-
-    assert model.ready
-    assert type(model._model) == TabularDrift
+    assert drift_detector.ready
+    assert type(drift_detector._model) == TabularDrift
 
 
 async def test_predict(
-    alibi_detect_tabular_drift_model: AlibiDetectRuntime,
+    drift_detector: AlibiDetectRuntime,
     inference_request: InferenceRequest,
 ):
-    response = await alibi_detect_tabular_drift_model.predict(inference_request)
-
+    # For #(batch - 1) requests, outputs should be empty
+    batch_size = drift_detector._ad_settings.batch_size
+    for _ in range(batch_size - 1):  # type: ignore
+        response = await drift_detector.predict(inference_request)
+        assert len(response.outputs) == 0
+
+    # By request batch_size, drift should run
+    response = await drift_detector.predict(inference_request)
     assert len(response.outputs) == 4
     assert response.outputs[0].name == "is_drift"
     assert response.outputs[0].shape == [1, 3]
@@ -78,13 +38,44 @@ async def test_predict(
     assert response.outputs[3].data[0] == P_VAL_THRESHOLD
 
 
-async def test_multiple_inputs_error(
-    alibi_detect_tabular_drift_model: AlibiDetectRuntime,
-    inference_request: InferenceRequest,
-):
-    inference_request.inputs.append(
-        RequestInput(name="input-1", shape=[1, 3], data=[[0, 1, 6]], datatype="FP32")
+async def test_predict_batch(drift_detector: AlibiDetectRuntime, mocker):
+    # Mock detector
+    mocked_predict = mocker.patch.object(drift_detector._model, "predict")
+
+    # For #(batch - 1) requests, outputs should be empty
+    batch_size = drift_detector._ad_settings.batch_size
+    expected = np.random.randint(10, size=(batch_size, 3))  # type: ignore
+    for idx in range(batch_size - 1):  # type: ignore
+        inference_request = NumpyRequestCodec.encode_request(expected[idx : idx + 1])
+        response = await drift_detector.predict(inference_request)
+
+        assert len(response.outputs) == 0
+        mocked_predict.assert_not_called()
+
+    inference_request = NumpyRequestCodec.encode_request(
+        expected[batch_size - 1 : batch_size]  # type: ignore
     )
+    await drift_detector.predict(inference_request)
+
+    mocked_predict.assert_called_once()
+    payload = mocked_predict.call_args.args[0]
+    np.testing.assert_array_equal(payload, expected)
+
 
-    with pytest.raises(CodecError):
-        await alibi_detect_tabular_drift_model.predict(inference_request)
+async def test_predict_batch_cleared(
+    drift_detector: AlibiDetectRuntime,
+    inference_request: InferenceRequest,
+):
+    # For #(batch - 1) requests, outputs should be empty
+    batch_size = drift_detector._ad_settings.batch_size
+    for _ in range(batch_size - 1):  # type: ignore
+        response = await drift_detector.predict(inference_request)
+        assert len(response.outputs) == 0
+
+    # By request batch_size, drift should run
+    response = await drift_detector.predict(inference_request)
+    assert len(response.outputs) > 0
+
+    # Batch should now be cleared (and started from scratch)
+    response = await drift_detector.predict(inference_request)
+    assert len(response.outputs) == 0
diff --git a/runtimes/alibi-detect/tests/test_outlier_detector.py b/runtimes/alibi-detect/tests/test_outlier_detector.py
--- a/runtimes/alibi-detect/tests/test_outlier_detector.py
+++ b/runtimes/alibi-detect/tests/test_outlier_detector.py
@@ -1,129 +1,23 @@
-import pytest
-from mlserver.settings import ModelSettings, ModelParameters
-from mlserver.types import RequestInput, InferenceRequest
-from mlserver.codecs import CodecError
-from mlserver_alibi_detect import AlibiDetectRuntime
 from alibi_detect.od import OutlierVAE
-from alibi_detect.utils.saving import save_detector
-import os
-import numpy as np
-import tensorflow as tf
-from tensorflow.keras.layers import Dense, InputLayer
-
-tf.keras.backend.clear_session()
-
-X_REF_THRESHOLD = 0.05
-
-
-@pytest.fixture
-def alibi_detect_tabular_outlier_model_settings(
-    alibi_detect_tabular_outlier_model_uri: str,
-) -> ModelSettings:
-    return ModelSettings(
-        name="alibi-detect-model",
-        implementation=AlibiDetectRuntime,
-        parameters=ModelParameters(
-            uri=alibi_detect_tabular_outlier_model_uri,
-            version="v1.2.3",
-            extra={
-                "predict_parameters": {
-                    "outlier_type": "instance",
-                    "return_feature_score": False,
-                    "return_instance_score": True,
-                }
-            },
-        ),
-    )
-
-
-@pytest.fixture
-def alibi_detect_tabular_outlier_model_uri(tmp_path) -> str:
-    X_ref = np.array([[1, 2, 3]])
-    n_features = X_ref.shape[1]
-    latent_dim = 2
-    encoder_net = tf.keras.Sequential(
-        [
-            InputLayer(input_shape=(n_features,)),
-            Dense(25, activation=tf.nn.relu),
-            Dense(10, activation=tf.nn.relu),
-            Dense(5, activation=tf.nn.relu),
-        ]
-    )
-
-    decoder_net = tf.keras.Sequential(
-        [
-            InputLayer(input_shape=(latent_dim,)),
-            Dense(5, activation=tf.nn.relu),
-            Dense(10, activation=tf.nn.relu),
-            Dense(25, activation=tf.nn.relu),
-            Dense(n_features, activation=None),
-        ]
-    )
-
-    od = OutlierVAE(
-        threshold=X_REF_THRESHOLD,
-        score_type="mse",
-        encoder_net=encoder_net,
-        decoder_net=decoder_net,
-        latent_dim=latent_dim,
-        samples=5,
-    )
-
-    od.fit(X_ref, loss_fn=tf.keras.losses.mse, epochs=5, verbose=True)
-
-    detector_uri = os.path.join(tmp_path, "alibi-detector-artifacts")
-    save_detector(od, detector_uri)
 
-    return detector_uri
+from mlserver.types import InferenceRequest
 
+from mlserver_alibi_detect import AlibiDetectRuntime
 
-@pytest.fixture
-async def alibi_detect_tabular_outlier_model(
-    alibi_detect_tabular_outlier_model_settings: ModelSettings,
-) -> AlibiDetectRuntime:
-    model = AlibiDetectRuntime(alibi_detect_tabular_outlier_model_settings)
-    await model.load()
-
-    return model
-
-
-async def test_load_folder(
-    alibi_detect_tabular_outlier_model_uri: str,
-    alibi_detect_tabular_outlier_model_settings: ModelSettings,
-):
-    alibi_detect_tabular_outlier_model_settings.parameters.uri = (  # type: ignore
-        alibi_detect_tabular_outlier_model_uri
-    )
-
-    model = AlibiDetectRuntime(alibi_detect_tabular_outlier_model_settings)
-    await model.load()
 
-    assert model.ready
-    assert type(model._model) == OutlierVAE
+async def test_load_folder(outlier_detector: AlibiDetectRuntime):
+    assert outlier_detector.ready
+    assert type(outlier_detector._model) == OutlierVAE
 
 
 async def test_predict(
-    alibi_detect_tabular_outlier_model: AlibiDetectRuntime,
+    outlier_detector: AlibiDetectRuntime,
     inference_request: InferenceRequest,
-    alibi_detect_tabular_outlier_model_settings: ModelSettings,
 ):
-    response = await alibi_detect_tabular_outlier_model.predict(inference_request)
+    response = await outlier_detector.predict(inference_request)
 
     assert len(response.outputs) == 3
     assert response.outputs[0].name == "instance_score"
     assert response.outputs[1].name == "feature_score"
     assert response.outputs[2].name == "is_outlier"
     assert response.outputs[2].shape == [1, 1]
-
-
-async def test_multiple_inputs_error(
-    alibi_detect_tabular_outlier_model: AlibiDetectRuntime,
-    inference_request: InferenceRequest,
-    alibi_detect_tabular_outlier_model_settings: ModelSettings,
-):
-    inference_request.inputs.append(
-        RequestInput(name="input-1", shape=[1, 3], data=[[0, 1, 6]], datatype="FP32")
-    )
-
-    with pytest.raises(CodecError):
-        await alibi_detect_tabular_outlier_model.predict(inference_request)
diff --git a/runtimes/alibi-detect/tests/test_runtime.py b/runtimes/alibi-detect/tests/test_runtime.py
new file mode 100644
--- /dev/null
+++ b/runtimes/alibi-detect/tests/test_runtime.py
@@ -0,0 +1,18 @@
+import pytest
+
+from mlserver.codecs import CodecError
+from mlserver.types import RequestInput, InferenceRequest
+
+from mlserver_alibi_detect import AlibiDetectRuntime
+
+
+async def test_multiple_inputs_error(
+    outlier_detector: AlibiDetectRuntime,
+    inference_request: InferenceRequest,
+):
+    inference_request.inputs.append(
+        RequestInput(name="input-1", shape=[1, 3], data=[[0, 1, 6]], datatype="FP32")
+    )
+
+    with pytest.raises(CodecError):
+        await outlier_detector.predict(inference_request)