Possible to use `params` argument of `mlflow.pyfunc.PythonModel` deployed to Databricks endpoint?

[dcunningham1]

I'm deploying a custom model using the `mlflow.pyfunc.PythonModel` class as described here. My model uses the `params` argument in the `predict` method to allow the user to choose some aspects of the model at inference time. For example:

class CustomModelWrapper(mlflow.pyfunc.PythonModel):
    def __init__(self, model):
        self.model = model

    def predict(self, context, model_input, params):
        confidence_level = params.get("confidence_level", 0.9)
        return self.model.predict(model_input, confidence_level=confidence_level)

The model works fine when I load it from the MLflow model registry. Passing in params works as expected.

But it does not work when I deploy the model to an endpoint on Databricks. It seems that Databricks does not support the `params` argument. That's unfortunate because it seems like it would be useful/important for many use cases.

Does Databricks have any plans to support the `params` argument in model serving endpoints in the future? In the meantime, is there any reasonable workaround to allow the user to specify parameters at inference time?

[umesh_gattem]

Is there any update on this feature?

[dcunningham1]

No updates that I'm aware of.

[KAdamatzky]

Hi, have you had any luck with this issue?

I am in the same position but I am not sure what the correct syntax is for passing params to the model endpoint via REST call.

[dcunningham1]

No luck. I was not able to pass params to the model endpoint.

I found a workaround that works for my use case, but it won't work for everyone. I changed the response format of my model. Specifically, I made my model return a dictionary where the keys represent different param values the user might be interested in. Instead of returning a single value (or a few values), the model returns a large dictionary of values.

In my use case, I want the model to return a point prediction plus an upper and lower bound. I wanted to use a parameter to allow the user to specify the size of the confidence interval (75%, 90%, etc.) using params at runtime. But instead, I made the model return a bunch of differently sized confidence intervals.

### I wanted it to work like this

class CustomModelWrapper(mlflow.pyfunc.PythonModel):
    def __init__(self, model):
        self.model = model

    def predict(self, context, model_input, params):
        confidence_level = params.get("confidence_level", 0.9)
        return self.model.predict(model_input, confidence_level=confidence_level)

# return value:
# {
#     "prediction": ..., # point prediction
#     "lower_bound": ..., # lower bound, depending on `confidence_level` param
#     "upper_bound": ..., # upper bound, depending on `confidence_level` param
# }


### but instead I made it work like this

class CustomModelWrapper(mlflow.pyfunc.PythonModel):
    def __init__(self, model):
        self.model = model

    def predict(self, context, model_input):
        # not handling any params here!!!
        return self.model.predict(model_input)

# return value:
# {
#     "prediction": ..., # point prediction
#     "interval_0.1": ..., # 10th percentile prediction
#     "interval_0.2": ..., # 20th percentile prediction
#     ..., # include all percentiles the user might need
#     "interval_0.9": ..., # 90th percentile prediction
# }
#
# (the user will have to select the keys that they're interested in)

(Note that the code for defining the return value format is all in my custom model code. It's obviously very dependent on the specific model and use case.)

This is not a great workaround! There are a few big problems:

1. You can't return every possible parameter value if your parameters are continuous. In my example, if the user wanted to get the 15th percentile prediction from the model, they couldn't do it.

2. Increased latency. Depending on your model and parameters, this workaround might make your API responses a lot slower because the model is doing more stuff. And most of that extra computation will be useless because the user doesn't need all of the output. This wasn't a concern for my use case, but it could be a big problem for many other use cases.

3. Not feasible if you have a lot of parameters. I only had a single parameter in my model, but this approach falls apart if you have more than one or two parameters because there will be too many possible combinations.