Load a pyfunc model logged with Feature Store

[SOlivero]

Hi, I'm using Databricks Feature Store to register a custom model using a model wrapper as follows:

# Log custom model to MLflow
fs.log_model(
  artifact_path="model",
  model = production_model,
  flavor = mlflow.pyfunc, 
  training_set = training_set,
  registered_model_name = model_name,
  conda_env=model_env
)

The model is correctly serving and I can score batch using the model's uri:

fs.score_batch('models:/model_name/latest',
                dataset,
                result_type = ArrayType(StringType()))

What I can not do is to load the model using: 

import mlflow
import databricks.feature_store
mlflow.pyfunc.load_model('models:/model_name/latest')

 I keep on getting the following error: 

ModuleNotFoundError: No module named 'databricks.feature_store.mlflow_model'

What is the correct way to load or log the model for this case? 

[Siebert_Looije]

Hi 

Thanks for your message. Databricks made an example with a notebook (with dummy data) available how you can access and use these cases. This might gives a direction where to look for: https://learn.microsoft.com/en-us/azure/databricks/_static/notebooks/machine-learning/feature-store-.... 

Hopefully this helps and else please send a message then I can have a look further.

[SOlivero]

Thank you for the help! 

I had already logged and score batch with my model. Nevertheless I can not do is to load the custom model object to another notebook with:

mlflow.pyfunc.load_model('models:/model_name/latest')

[Kumaran]

Hi @SOlivero 

Make sure that the model was in fact saved with the provided URI.

The latest keyword will retrieve the latest version of the registered model when mlflow.pyfunc.load_model('models:/model_name/latest') is executed, not the highest version. If you have not registered the model with the name model_name in the models registry you can try to load it using the full path:

mlflow.pyfunc.load_model("file:///path/to/the/saved/model")

where /path/to/the/saved/model represents the location where the model artifacts were saved during training (i.e the path where mlflow.pyfunc.save_model() was called during training).

If this doesn't help, check that you can access the model outside of the notebook. You could also check that the Databricks cluster you're running on has access to the Databricks models registry if that is where your model was saved.