Sherif akoush/mlsv 42/add alibi explain runtime

mlserver/codecs/numpy.py

@@ -93,16 +93,16 @@ def encode(cls, name: str, payload: np.ndarray) -> ResponseOutput:
         )
 
     @classmethod
-    def decode(cls, v2_data: RequestInput) -> np.ndarray:
-        model_data = _to_ndarray(v2_data)
+    def decode(cls, request_input: RequestInput) -> np.ndarray:
+        model_data = _to_ndarray(request_input)
 
         # TODO: Check if reshape not valid
-        return model_data.reshape(v2_data.shape)
+        return model_data.reshape(request_input.shape)
 
     @classmethod
-    def decode_response_output(cls, v2_data: ResponseOutput) -> np.ndarray:
+    def decode_response_output(cls, response_output: ResponseOutput) -> np.ndarray:
         # TODO: merge this logic with `decode`
-        return cls.decode(v2_data)  # type: ignore
+        return cls.decode(response_output)  # type: ignore
 
     @classmethod
     def encode_request_input(cls, name: str, payload: np.ndarray) -> RequestInput:

mlserver/codecs/string.py

@@ -54,12 +54,13 @@ def decode(cls, request_input: RequestInput) -> List[str]:
     @classmethod
     def encode_request_input(cls, name: str, payload: List[str]) -> RequestInput:
         # TODO: merge this logic with `encode`
-        data = cls.encode(name=name, payload=payload)
+        # note: this will only work with REST and not grpc as we might have
+        # variable length strings
         return RequestInput(
-            name=data.name,
-            datatype=data.datatype,
-            shape=data.shape,
-            data=data.data.__root__.decode("ascii"),  # to allow json serialisation
+            name=name,
+            datatype="BYTES",
+            shape=[len(payload), -1],  # this is discarded downstream?
+            data=payload,
             parameters=Parameters(content_type=cls.ContentType),
         )
 

mlserver/errors.py

@@ -24,3 +24,13 @@ def __init__(self, name: str, version: str = None):
 class InferenceError(MLServerError):
     def __init__(self, msg: str):
         super().__init__(msg)
+
+
+class RemoteInferenceError(MLServerError):
+    def __init__(self, code: int, reason: str):
+        super().__init__(f"Remote inference call failed with {code}, {reason}")
+
+
+class ModelParametersMissing(MLServerError):
+    def __init__(self, model_name: str):
+        super().__init__(f"Parameters missing for model {model_name}")

runtimes/alibi-detect/mlserver_alibi_detect/runtime.py

@@ -46,9 +46,11 @@ class AlibiDetectRuntime(MLModel):
     """
 
     def __init__(self, settings: ModelSettings):
-
-        extra = settings.parameters.extra  # type: ignore
-        self.alibi_detect_settings = AlibiDetectSettings(**extra)  # type: ignore
+        if settings.parameters is None:
+            self.alibi_detect_settings = AlibiDetectSettings()
+        else:
+            extra = settings.parameters.extra
+            self.alibi_detect_settings = AlibiDetectSettings(**extra)  # type: ignore
         super().__init__(settings)
 
     @custom_handler(rest_path="/")

runtimes/alibi-explain/mlserver_alibi_explain/alibi_dependency_reference.py

@@ -0,0 +1,62 @@
+from dataclasses import dataclass
+from enum import Enum
+from typing import Union, Dict
+
+
+@dataclass
+class ExplainerDependencyReference:
+    """Class for keeping track of dependencies required to Alibi runtime."""
+
+    explainer_name: str
+    alibi_class: str
+    runtime_class: str
+
+
+_ANCHOR_IMAGE_TAG = "anchor_image"
+_ANCHOR_TEXT_TAG = "anchor_text"
+_INTEGRATED_GRADIENTS_TAG = "integrated_gradients"
+
+
+# NOTE: to add new explainers populate the below dict with a new
+# ExplainerDependencyReference, referencing the specific runtime class in mlserver
+# and the specific alibi explain class.
+# this can be simplified when alibi moves to a config based init.
+
+_TAG_TO_RT_IMPL: Dict[str, ExplainerDependencyReference] = {
+    _ANCHOR_IMAGE_TAG: ExplainerDependencyReference(
+        explainer_name=_ANCHOR_IMAGE_TAG,
+        runtime_class="mlserver_alibi_explain.explainers.black_box_runtime."
+        "AlibiExplainBlackBoxRuntime",
+        alibi_class="alibi.explainers.AnchorImage",
+    ),
+    _ANCHOR_TEXT_TAG: ExplainerDependencyReference(
+        explainer_name=_ANCHOR_TEXT_TAG,
+        runtime_class="mlserver_alibi_explain.explainers.black_box_runtime."
+        "AlibiExplainBlackBoxRuntime",
+        alibi_class="alibi.explainers.AnchorText",
+    ),
+    _INTEGRATED_GRADIENTS_TAG: ExplainerDependencyReference(
+        explainer_name=_INTEGRATED_GRADIENTS_TAG,
+        runtime_class="mlserver_alibi_explain.explainers.integrated_gradients."
+        "IntegratedGradientsWrapper",
+        alibi_class="alibi.explainers.IntegratedGradients",
+    ),
+}
+
+
+class ExplainerEnum(str, Enum):
+    anchor_image = _ANCHOR_IMAGE_TAG
+    anchor_text = _ANCHOR_TEXT_TAG
+    integrated_gradients = _INTEGRATED_GRADIENTS_TAG
+
+
+def get_mlmodel_class_as_str(tag: Union[ExplainerEnum, str]) -> str:
+    if isinstance(tag, ExplainerEnum):
+        tag = tag.value
+    return _TAG_TO_RT_IMPL[tag].runtime_class
+
+
+def get_alibi_class_as_str(tag: Union[ExplainerEnum, str]) -> str:
+    if isinstance(tag, ExplainerEnum):
+        tag = tag.value
+    return _TAG_TO_RT_IMPL[tag].alibi_class

runtimes/alibi-explain/mlserver_alibi_explain/common.py

@@ -2,7 +2,6 @@
 import contextvars
 import functools
 from asyncio import AbstractEventLoop
-from enum import Enum
 from importlib import import_module
 from typing import Any, Optional, Type, Callable, Awaitable, Union, List
 
@@ -11,6 +10,7 @@
 from pydantic import BaseSettings
 
 from mlserver.codecs import StringCodec, NumpyCodec
+from mlserver.errors import RemoteInferenceError
 from mlserver.types import (
     ResponseOutput,
     InferenceResponse,
@@ -22,39 +22,10 @@
 
 _MAX_RETRY_ATTEMPT = 3
 
-_ANCHOR_IMAGE_TAG = "anchor_image"
-_ANCHOR_TEXT_TAG = "anchor_text"
-_INTEGRATED_GRADIENTS_TAG = "integrated_gradients"
-
-_TAG_TO_RT_IMPL = {
-    _ANCHOR_IMAGE_TAG: (
-        "mlserver_alibi_explain.explainers.black_box_runtime."
-        "AlibiExplainBlackBoxRuntime",
-        "alibi.explainers.AnchorImage",
-    ),
-    _ANCHOR_TEXT_TAG: (
-        "mlserver_alibi_explain.explainers.black_box_runtime."
-        "AlibiExplainBlackBoxRuntime",
-        "alibi.explainers.AnchorText",
-    ),
-    _INTEGRATED_GRADIENTS_TAG: (
-        "mlserver_alibi_explain.explainers.integrated_gradients."
-        "IntegratedGradientsWrapper",
-        "alibi.explainers.IntegratedGradients",
-    ),
-}
-
-
 ENV_PREFIX_ALIBI_EXPLAIN_SETTINGS = "MLSERVER_MODEL_ALIBI_EXPLAIN_"
 EXPLAIN_PARAMETERS_TAG = "explain_parameters"
 
 
-class ExplainerEnum(str, Enum):
-    anchor_image = _ANCHOR_IMAGE_TAG
-    anchor_text = _ANCHOR_TEXT_TAG
-    integrated_gradients = _INTEGRATED_GRADIENTS_TAG
-
-
 def convert_from_bytes(output: ResponseOutput, ty: Optional[Type]) -> Any:
     """
     This utility function decodes the response from bytes string to python object dict.
@@ -76,7 +47,7 @@ def remote_predict(
 ) -> InferenceResponse:
     response_raw = requests.post(predictor_url, json=v2_payload.dict())
     if response_raw.status_code != 200:
-        raise ValueError(f"{response_raw.status_code} / {response_raw.reason}")
+        raise RemoteInferenceError(response_raw.status_code, response_raw.reason)
     return InferenceResponse.parse_raw(response_raw.text)
 
 
@@ -92,18 +63,6 @@ def execute_async(
     return loop.run_in_executor(None, func_call)
 
 
-def get_mlmodel_class_as_str(tag: Union[ExplainerEnum, str]) -> str:
-    if isinstance(tag, ExplainerEnum):
-        tag = tag.value
-    return _TAG_TO_RT_IMPL[tag][0]
-
-
-def get_alibi_class_as_str(tag: Union[ExplainerEnum, str]) -> str:
-    if isinstance(tag, ExplainerEnum):
-        tag = tag.value
-    return _TAG_TO_RT_IMPL[tag][1]
-
-
 class AlibiExplainSettings(BaseSettings):
     """
     Parameters that apply only to alibi explain models
@@ -112,7 +71,7 @@ class AlibiExplainSettings(BaseSettings):
     class Config:
         env_prefix = ENV_PREFIX_ALIBI_EXPLAIN_SETTINGS
 
-    infer_uri: Optional[str]
+    infer_uri: str
     explainer_type: str
     init_parameters: Optional[dict]
 

runtimes/alibi-explain/mlserver_alibi_explain/explainers/black_box_runtime.py

@@ -1,12 +1,10 @@
-from typing import Type, Any, Dict, Optional, List, Union
+from typing import Type, Any, Dict, List, Union
 
 import numpy as np
 from alibi.api.interfaces import Explanation, Explainer
-from alibi.saving import load_explainer
 
 from mlserver import ModelSettings
 from mlserver.codecs import NumpyCodec
-from mlserver.settings import ModelParameters
 from mlserver_alibi_explain.common import (
     AlibiExplainSettings,
     remote_predict,
@@ -24,9 +22,14 @@ class AlibiExplainBlackBoxRuntime(AlibiExplainRuntimeBase):
     def __init__(self, settings: ModelSettings, explainer_class: Type[Explainer]):
         self._explainer_class = explainer_class
 
-        extra = settings.parameters.extra  # type: ignore
+        # if we are here we are sure that settings.parameters is set,
+        # just helping mypy
+        assert settings.parameters is not None
+        extra = settings.parameters.extra
         explainer_settings = AlibiExplainSettings(**extra)  # type: ignore
 
+        self.infer_uri = explainer_settings.infer_uri
+
         # TODO: validate the settings are ok with this specific explainer
         super().__init__(settings, explainer_settings)
 
@@ -37,12 +40,7 @@ async def load(self) -> bool:
             init_parameters["predictor"] = self._infer_impl
             self._model = self._explainer_class(**init_parameters)  # type: ignore
         else:
-            # load the model from disk
-            model_parameters: Optional[ModelParameters] = self.settings.parameters
-            assert model_parameters is not None
-            uri = model_parameters.uri  # type ignore
-            assert uri is not None, "uri has to be set"
-            self._model = load_explainer(uri, predictor=self._infer_impl)
+            self._model = self._load_from_uri(self._infer_impl)
 
         self.ready = True
         return self.ready
@@ -63,10 +61,9 @@ def _infer_impl(self, input_data: Union[np.ndarray, List[str]]) -> np.ndarray:
 
         v2_request = to_v2_inference_request(input_data)
 
-        # TODO add some exception handling here
-        infer_uri = self.alibi_explain_settings.infer_uri
-        assert infer_uri is not None, "infer_uri has to be set"
-        v2_response = remote_predict(v2_payload=v2_request, predictor_url=infer_uri)
+        v2_response = remote_predict(
+            v2_payload=v2_request, predictor_url=self.infer_uri
+        )
 
         # TODO: do we care about more than one output?
         return NumpyCodec.decode_response_output(v2_response.outputs[0])

runtimes/alibi-explain/mlserver_alibi_explain/explainers/integrated_gradients.py

@@ -3,6 +3,7 @@
 import tensorflow as tf
 from alibi.api.interfaces import Explanation
 
+from mlserver.errors import InvalidModelURI
 from mlserver_alibi_explain.explainers.white_box_runtime import (
     AlibiExplainWhiteBoxRuntime,
 )
@@ -17,4 +18,9 @@ def _explain_impl(self, input_data: Any, explain_parameters: Dict) -> Explanatio
 
     async def _get_inference_model(self) -> Any:
         inference_model_path = self.alibi_explain_settings.infer_uri
-        return tf.keras.models.load_model(inference_model_path)
+        try:
+            model = tf.keras.models.load_model(inference_model_path)
+        except IOError:
+            raise InvalidModelURI(self.name, inference_model_path)
+
+        return model

runtimes/alibi-explain/mlserver_alibi_explain/explainers/white_box_runtime.py

@@ -1,11 +1,9 @@
 from abc import ABC
-from typing import Any, Type, Optional
+from typing import Any, Type
 
 from alibi.api.interfaces import Explainer
-from alibi.saving import load_explainer
 
 from mlserver import ModelSettings
-from mlserver.settings import ModelParameters
 from mlserver_alibi_explain.common import AlibiExplainSettings
 from mlserver_alibi_explain.runtime import AlibiExplainRuntimeBase
 
@@ -21,7 +19,10 @@ def __init__(self, settings: ModelSettings, explainer_class: Type[Explainer]):
         self._inference_model = None
         self._explainer_class = explainer_class
 
-        extra = settings.parameters.extra  # type: ignore
+        # if we are here we are sure that settings.parameters is set,
+        # just helping mypy
+        assert settings.parameters is not None
+        extra = settings.parameters.extra
         explainer_settings = AlibiExplainSettings(**extra)  # type: ignore
 
         # TODO: validate the settings are ok with this specific explainer
@@ -36,13 +37,7 @@ async def load(self) -> bool:
             init_parameters["model"] = self._inference_model
             self._model = self._explainer_class(**init_parameters)  # type: ignore
         else:
-            # load the model from disk
-            # full model is passed as `predictor`
-            model_parameters: Optional[ModelParameters] = self.settings.parameters
-            assert model_parameters is not None
-            uri = model_parameters.uri  # type ignore
-            assert uri is not None, "uri has to be set"
-            self._model = load_explainer(uri, predictor=self._inference_model)
+            self._model = self._load_from_uri(self._inference_model)
 
         self.ready = True
         return self.ready

runtimes/alibi-explain/mlserver_alibi_explain/runtime.py

@@ -1,10 +1,12 @@
 from typing import Any, Optional, List, Dict
 
-from alibi.api.interfaces import Explanation
+from alibi.api.interfaces import Explanation, Explainer
+from alibi.saving import load_explainer
 
 from mlserver.codecs import NumpyCodec, InputCodec, StringCodec
+from mlserver.errors import ModelParametersMissing, InvalidModelURI
 from mlserver.model import MLModel
-from mlserver.settings import ModelSettings
+from mlserver.settings import ModelSettings, ModelParameters
 from mlserver.types import (
     InferenceRequest,
     InferenceResponse,
@@ -17,12 +19,14 @@
 from mlserver_alibi_explain.common import (
     execute_async,
     AlibiExplainSettings,
-    get_mlmodel_class_as_str,
-    get_alibi_class_as_str,
     import_and_get_class,
     EXPLAIN_PARAMETERS_TAG,
     EXPLAINER_TYPE_TAG,
 )
+from mlserver_alibi_explain.alibi_dependency_reference import (
+    get_mlmodel_class_as_str,
+    get_alibi_class_as_str,
+)
 
 
 class AlibiExplainRuntimeBase(MLModel):
@@ -72,7 +76,19 @@ async def _async_explain_impl(
             explain_parameters=explain_parameters,
         )
         # TODO: Convert alibi-explain output to v2 protocol, for now we use to_json
-        return StringCodec.encode(payload=[explanation.to_json()], name="explain")
+        return StringCodec.encode(payload=[explanation.to_json()], name="explanation")
+
+    def _load_from_uri(self, predictor: Any) -> Explainer:
+        # load the model from disk
+        # full model is passed as `predictor`
+        # load the model from disk
+        model_parameters: Optional[ModelParameters] = self.settings.parameters
+        if model_parameters is None:
+            raise ModelParametersMissing(self.name)
+        uri = model_parameters.uri
+        if uri is None:
+            raise InvalidModelURI(self.name)
+        return load_explainer(uri, predictor=predictor)
 
     def _explain_impl(self, input_data: Any, explain_parameters: Dict) -> Explanation:
         """Actual explain to be implemented by subclasses"""