MLFlow model with Array type in signature not loading

[wbinek]

I was trying to serve a model in MLFlow Model format created using MLFlow > 2.10. The MLFLow 2.10 added support for Array and Object datatypes in signature (https://www.mlflow.org/docs/latest/model/signatures.html), which seems to be unsupported by the MLServer.

Error:

2024-06-12 08:10:45,410 [mlserver.parallel] ERROR - An error occurred processing a model update of type 'Load'.
Traceback (most recent call last):
  File "/opt/conda/lib/python3.10/site-packages/mlserver/parallel/worker.py", line 158, in _process_model_update
    await self._model_registry.load(model_settings)
  File "/opt/conda/lib/python3.10/site-packages/mlserver/registry.py", line 293, in load
    return await self._models[model_settings.name].load(model_settings)
  File "/opt/conda/lib/python3.10/site-packages/mlserver/registry.py", line 148, in load
    await self._load_model(new_model)
  File "/opt/conda/lib/python3.10/site-packages/mlserver/registry.py", line 165, in _load_model
    model.ready = await model.load()
  File "/opt/conda/lib/python3.10/site-packages/mlserver_mlflow/runtime.py", line 159, in load
    self._sync_metadata()
  File "/opt/conda/lib/python3.10/site-packages/mlserver_mlflow/runtime.py", line 171, in _sync_metadata
    self.inputs = to_metadata_tensors(
  File "/opt/conda/lib/python3.10/site-packages/mlserver_mlflow/metadata.py", line 57, in to_metadata_tensors
    datatype, content_type = _get_content_type(input_spec)
  File "/opt/conda/lib/python3.10/site-packages/mlserver_mlflow/metadata.py", line 41, in _get_content_type
    return _MLflowToContentType[input_spec.type]
TypeError: unhashable type: 'Array'
2024-06-12 08:10:45,411 [mlserver] INFO - Couldn't load model 'mlflow-model'. Model will be removed from registry.

MLFlow Model Signature:

signature:
  inputs: '[{"type": "datetime", "name": "start", "required": true}, 
            {"type": "array", "items": {"type": "array", "items": {"type": "long"}}, "name": "target", "required": true}, 
            {"type": "array", "items": {"type": "string"}, "name": "feat_static_cat", "required": true}, 
            {"type": "array", "items": {"type": "array", "items": {"type": "boolean"}}, "name": "feat_dynamic_real", "required": true}]'
  outputs: null
  params: null

[lurecas]

Hello there!

I also stumbled upon this issue and after digging a little bit, I did a small modification on the get_content_type method from the metadata.py file (that belongs to the MLflow Runtime):

def _get_content_type(input_spec: InputSpec) -> Tuple[MDatatype, str]:
    if isinstance(input_spec, TensorSpec):
        datatype = to_datatype(input_spec.type)
        content_type = NumpyCodec.ContentType
        return datatype, content_type
    elif isinstance(input_spec, DataType):
        return _MLflowToContentType[input_spec]
    elif isinstance(input_spec.type, Array):
        return (MDatatype.BYTES, NumpyCodec.ContentType)
    elif isinstance(input_spec.type, Object):
        return (MDatatype.BYTES, NumpyCodec.ContentType)
    else:
        # TODO: Check if new type, which may not exist
        return _MLflowToContentType[input_spec.type]

For the moment, this has allowed me to run a couple of models that have object and array in their signatures. I will be opening a Pull Request with this addition (to be honest, I will be welcoming any kind of feedback, as I don't if that's the right codec to use for those cases), but I wanted to share this in case anyone can benefit from it.

Cheers!