[BUG] Loading Pyfunc Class missing custom instance methods

[tallen94]

Willingness to contribute

The MLflow Community encourages bug fix contributions. Would you or another member of your organization be willing to contribute a fix for this bug to the MLflow code base?

• [ ] Yes. I can contribute a fix for this bug independently.
• [x] Yes. I would be willing to contribute a fix for this bug with guidance from the MLflow community.
• [ ] No. I cannot contribute a bug fix at this time.

System information

• Have I written custom code (as opposed to using a stock example script provided in MLflow): Yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu
• MLflow installed from (source or binary): binary
• MLflow version (run mlflow --version): 1.11.0
• Python version: 3.6.8
• npm version, if running the dev UI:
• Exact command to reproduce: calling mlflow.pyfunc.load_model()

Describe the problem

Describe the problem clearly here. Include descriptions of the expected behavior and the actual behavior.

Im not sure if this was introduced as a feature or not, but here is the description:

I have a python function model that wraps around an sklearn model, a class that looks like this:

class GenericModel(mlflow.pyfunc.PythonModel):

  def __init__(self, model):
    self.model = model

  def predict(self, input):
    return self.model.predict_proba(input)[:, 1][0]

  def get_model(self):
    return self.model

My _load_pyfunc() function returns an instance of this class above.

In version mlflow==1.7.2 when I call mlflow.pyfunc.load_model I get back the correct instance and I can call get_model() on the returned object. However in recent updates, like mlflow==1.11.0, this method is no longer available on the class.

I looked through the code and it seems that all that is happening is calling this _load_pyfunc() method which should return an instance of this class, but it seems that something is strongly typing it and removing the other custom instance methods I have written on it.

Am I just missing some information on how python classes work or how the import_module is working? Not entirely sure if this is an feature with MLFlow or just some python stuff.

Thank you for any help!!

Code to reproduce issue

Provide a reproducible test case that is the bare minimum necessary to generate the problem.

Create a generic PyfuncModel with additional instance methods and see if they are available when calling mlflow.pyfunc.load_model()

What component(s), interfaces, languages, and integrations does this bug affect?

Components

• [ ] area/artifacts: Artifact stores and artifact logging
• [ ] area/build: Build and test infrastructure for MLflow
• [ ] area/docs: MLflow documentation pages
• [ ] area/examples: Example code
• [ ] area/model-registry: Model Registry service, APIs, and the fluent client calls for Model Registry
• [x] area/models: MLmodel format, model serialization/deserialization, flavors
• [ ] area/projects: MLproject format, project running backends
• [ ] area/scoring: Local serving, model deployment tools, spark UDFs
• [ ] area/server-infra: MLflow server, JavaScript dev server
• [ ] area/tracking: Tracking Service, tracking client APIs, autologging

[vibss2397]

Hey, I am facing this issue as well, but in my case even the downgrading to mlflow 1.7.2 is not working. Any ideas on what to do?

[SoulEvill]

I am facing the same issue and am not sure if any of you guys were able to figure it out. it is very annoying and I was expecting to get some custom function to get model features all that stuff. please help, thank you!

[SoulEvill]

I think I found the reason the load model will call _load_pyfunc which will instantiate _PythonModelPyfuncWrapper class that only contains predict function. all customized functions will NOT pass to this. it is not flexible at all. not sure why it was implemented this way.

Not sure what would be the best way to fix it, any ideas?

code link: https://www.mlflow.org/docs/latest/_modules/mlflow/pyfunc/model.html

[nickhill97]

I'm experiencing this issue also, is there any update on a solution?

[vibss2397]

so i just use the method : model_class_instance._model_impl.python_model, this brings back the python class and all the methods associated with the original implementation. Hope it helps

[Jason-CKY]

thanks for this workaround! Wonder if it is possible to just expose these custom functions with the PyFuncModel class instead.