pythondatabricksmlflow

Return the list of all registered models in MLflow


I'm using mlflow on databricks and I've trained some models, that I can see in the "registered models" page.

Is there a way to extract the list of these models in code?

Something like

import mlflow
mlflow.set_tracking_uri("databricks")
model_infos = mlflow.tracking.MlflowClient().some_method_to_list_registered_models() 

for model_info in model_infos:
    print(f"Name: {model_info.name}, Version: {model_info.latest_versions[0].version}")

Solution

  • To get all registered models

    mlflow.search_registered_models()
    

    You will have to write a proper function to convert it into the required format here is an example for you:

    def get_Models():
        client = mlflow.MlflowClient()
        data = client.search_registered_models()
        models = []
        for model in data:
            models.append(model.name)
        result = []
        for model in models:
            model_versions = {"name": model}
            data = client.search_model_versions(filter_string =f"name='{model}'", order_by=["version_number DESC"])
            versions = list(map(lambda x: dict(x), data))
            model_versions["latest_versions"] = versions
            result.append(model_versions)
        return result