Repository: mlflow/mlflow

Issue:
[FR] Add Python fluent API for search_registered_models
### Willingness to contribute

No. I cannot contribute this feature at this time.

### Proposal Summary

Introduce a Python API for search_registered_models:

```
mlflow.search_registered_models(filter_string: Optional[str] = None, max_results: int = 100, order_by: Optional[List[str]] = None)
```

### Motivation

> #### What is the use case for this feature?

Help users easily search registered models without having to handle pagination

> #### Why is this use case valuable to support for MLflow users in general? ^ 

> #### Why is this use case valuable to support for your project(s) or organization? ^ 

> #### Why is it currently difficult to achieve this use case? Users must rely on the https://mlflow.org/docs/latest/python_api/mlflow.client.html#mlflow.client.MlflowClient.search_registered_models API, which requires them to handle page tokens manually


### Details

The API should be implemented in https://github.com/mlflow/mlflow/blob/7f2152370b8e93ed7bf8cc3e5e1e7ff39cd69ca5/mlflow/tracking/fluent.py

The only difference between this API and https://mlflow.org/docs/latest/python_api/mlflow.client.html#mlflow.client.MlflowClient.search_registered_models is that it should automatically handle pagination: instead of including a `page_token` argument, the API should 1. call https://mlflow.org/docs/latest/python_api/mlflow.client.html#mlflow.client.MlflowClient.search_registered_models, 2. retrieve the page token from the PagedList, 3. Call https://mlflow.org/docs/latest/python_api/mlflow.client.html#mlflow.client.MlflowClient.search_registered_models with the page token and append the results to the previous PagedList, 4. Continue until https://mlflow.org/docs/latest/python_api/mlflow.client.html#mlflow.client.MlflowClient.search_registered_models no longer returns a page token.

For reference, see https://github.com/mlflow/mlflow/blob/7f2152370b8e93ed7bf8cc3e5e1e7ff39cd69ca5/mlflow/tracking/fluent.py#L1154-L1167. MLflow provides a utility to help with pagination: https://github.com/mlflow/mlflow/blob/7f2152370b8e93ed7bf8cc3e5e1e7ff39cd69ca5/mlflow/utils/__init__.py#L224.

We can add some simple tests to https://github.com/mlflow/mlflow/blob/master/tests/tracking/fluent/test_fluent.py.

### What component(s) does this bug affect?

- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [X] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/recipes`: Recipes, Recipe APIs, Recipe configs, Recipe Templates
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [ ] `area/server-infra`: MLflow Tracking server backend
- [ ] `area/tracking`: Tracking Service, tracking client APIs, autologging

### What interface(s) does this bug affect?

- [ ] `area/uiux`: Front-end, user experience, plotting, JavaScript, JavaScript dev server
- [ ] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

### What language(s) does this bug affect?

- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

### What integration(s) does this bug affect?

- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations


Discussion:
I'll try this one
> I'll try this one

Thank you, @TSienki ! I've assigned the issue to you. Please let me know if you have any questions!

Resolving patch:
diff --git a/mlflow/__init__.py b/mlflow/__init__.py
--- a/mlflow/__init__.py
+++ b/mlflow/__init__.py
@@ -144,6 +144,7 @@
     get_experiment,
     get_experiment_by_name,
     search_experiments,
+    search_registered_models,
     create_experiment,
     set_experiment,
     log_params,
@@ -198,6 +199,7 @@
     "get_experiment",
     "get_experiment_by_name",
     "search_experiments",
+    "search_registered_models",
     "create_experiment",
     "set_experiment",
     "delete_experiment",
diff --git a/mlflow/tracking/client.py b/mlflow/tracking/client.py
--- a/mlflow/tracking/client.py
+++ b/mlflow/tracking/client.py
@@ -1924,9 +1924,10 @@ def search_registered_models(
             # prefix pattern
             filter_string = "name LIKE 'Boston%'"
             results = client.search_registered_models(filter_string=filter_string)
+            print("-" * 80)
             for res in results:
                 for mv in res.latest_versions:
-                print("name={}; run_id={}; version={}".format(mv.name, mv.run_id, mv.version))
+                    print("name={}; run_id={}; version={}".format(mv.name, mv.run_id, mv.version))
 
             # Get all registered models and order them by ascending order of the names
             results = client.search_registered_models(order_by=["name ASC"])
diff --git a/mlflow/tracking/fluent.py b/mlflow/tracking/fluent.py
--- a/mlflow/tracking/fluent.py
+++ b/mlflow/tracking/fluent.py
@@ -12,11 +12,13 @@
 
 from mlflow.entities import Experiment, Run, RunStatus, Param, RunTag, Metric, ViewType
 from mlflow.entities.lifecycle_stage import LifecycleStage
+from mlflow.entities.model_registry import RegisteredModel
 from mlflow.exceptions import MlflowException
 from mlflow.protos.databricks_pb2 import (
     INVALID_PARAMETER_VALUE,
     RESOURCE_DOES_NOT_EXIST,
 )
+from mlflow.store.model_registry import SEARCH_REGISTERED_MODEL_MAX_RESULTS_DEFAULT
 from mlflow.tracking.client import MlflowClient
 from mlflow.tracking import artifact_utils, _get_store
 from mlflow.tracking.context import registry as context_registry
@@ -1167,6 +1169,110 @@ def pagination_wrapper_func(number_to_get, next_page_token):
     )
 
 
+def search_registered_models(
+    max_results: Optional[int] = None,
+    filter_string: Optional[str] = None,
+    order_by: Optional[List[str]] = None,
+) -> List[RegisteredModel]:
+    """
+    Search for registered models that satisfy the filter criteria.
+
+    :param filter_string: Filter query string
+        (e.g., ``"name = 'a_model_name' and tag.key = 'value1'"``),
+        defaults to searching for all registered models. The following identifiers, comparators,
+        and logical operators are supported.
+
+        Identifiers
+          - ``name``: registered model name.
+          - ``tags.<tag_key>``: registered model tag. If ``tag_key`` contains spaces, it must be
+            wrapped with backticks (e.g., ``"tags.`extra key`"``).
+
+        Comparators
+          - ``=``: Equal to.
+          - ``!=``: Not equal to.
+          - ``LIKE``: Case-sensitive pattern match.
+          - ``ILIKE``: Case-insensitive pattern match.
+
+        Logical operators
+          - ``AND``: Combines two sub-queries and returns True if both of them are True.
+
+    :param max_results: If passed, specifies the maximum number of models desired. If not
+                        passed, all models will be returned.
+    :param order_by: List of column names with ASC|DESC annotation, to be used for ordering
+                     matching search results.
+    :return: A list of :py:class:`mlflow.entities.model_registry.RegisteredModel` objects
+             that satisfy the search expressions.
+
+    .. code-block:: python
+        :caption: Example
+
+        import mlflow
+        from sklearn.linear_model import LogisticRegression
+
+        with mlflow.start_run():
+            mlflow.sklearn.log_model(
+                LogisticRegression(),
+                "Cordoba",
+                registered_model_name="CordobaWeatherForecastModel",
+            )
+            mlflow.sklearn.log_model(
+                LogisticRegression(),
+                "Boston",
+                registered_model_name="BostonWeatherForecastModel",
+            )
+
+        # Get search results filtered by the registered model name
+        filter_string = "name = 'CordobaWeatherForecastModel'"
+        results = mlflow.search_registered_models(filter_string=filter_string)
+        print("-" * 80)
+        for res in results:
+            for mv in res.latest_versions:
+                print("name={}; run_id={}; version={}".format(mv.name, mv.run_id, mv.version))
+
+        # Get search results filtered by the registered model name that matches
+        # prefix pattern
+        filter_string = "name LIKE 'Boston%'"
+        results = mlflow.search_registered_models(filter_string=filter_string)
+        print("-" * 80)
+        for res in results:
+            for mv in res.latest_versions:
+                print("name={}; run_id={}; version={}".format(mv.name, mv.run_id, mv.version))
+
+        # Get all registered models and order them by ascending order of the names
+        results = mlflow.search_registered_models(order_by=["name ASC"])
+        print("-" * 80)
+        for res in results:
+            for mv in res.latest_versions:
+                print("name={}; run_id={}; version={}".format(mv.name, mv.run_id, mv.version))
+
+    .. code-block:: text
+        :caption: Output
+
+        --------------------------------------------------------------------------------
+        name=CordobaWeatherForecastModel; run_id=248c66a666744b4887bdeb2f9cf7f1c6; version=1
+        --------------------------------------------------------------------------------
+        name=BostonWeatherForecastModel; run_id=248c66a666744b4887bdeb2f9cf7f1c6; version=1
+        --------------------------------------------------------------------------------
+        name=BostonWeatherForecastModel; run_id=248c66a666744b4887bdeb2f9cf7f1c6; version=1
+        name=CordobaWeatherForecastModel; run_id=248c66a666744b4887bdeb2f9cf7f1c6; version=1
+
+    """
+
+    def pagination_wrapper_func(number_to_get, next_page_token):
+        return MlflowClient().search_registered_models(
+            max_results=number_to_get,
+            filter_string=filter_string,
+            order_by=order_by,
+            page_token=next_page_token,
+        )
+
+    return get_results_from_paginated_fn(
+        pagination_wrapper_func,
+        SEARCH_REGISTERED_MODEL_MAX_RESULTS_DEFAULT,
+        max_results,
+    )
+
+
 def create_experiment(
     name: str,
     artifact_location: Optional[str] = None,


Test patch:
diff --git a/tests/tracking/fluent/test_fluent.py b/tests/tracking/fluent/test_fluent.py
--- a/tests/tracking/fluent/test_fluent.py
+++ b/tests/tracking/fluent/test_fluent.py
@@ -1,5 +1,8 @@
 from collections import defaultdict
 from importlib import reload
+from itertools import zip_longest
+
+from mlflow.store.model_registry import SEARCH_REGISTERED_MODEL_MAX_RESULTS_DEFAULT
 from mlflow.store.tracking import SEARCH_MAX_RESULTS_DEFAULT
 
 import os
@@ -368,12 +371,8 @@ def test_search_experiments(tmp_path):
 
     active_experiment_names = [f"active_{i}" for i in range(num_active_experiments)]
     tag_values = ["x", "x", "y"]
-    for (idx, active_experiment_name) in enumerate(active_experiment_names):
-        if idx < 3:
-            tags = {"tag": tag_values[idx]}
-        else:
-            tags = None
-        mlflow.create_experiment(active_experiment_name, tags=tags)
+    for (tag, active_experiment_name) in zip_longest(tag_values, active_experiment_names):
+        mlflow.create_experiment(active_experiment_name, tags={"tag": tag} if tag else None)
 
     deleted_experiment_names = [f"deleted_{i}" for i in range(num_deleted_experiments)]
     for deleted_experiment_name in deleted_experiment_names:
@@ -423,6 +422,54 @@ def test_search_experiments(tmp_path):
     assert [e.name for e in experiments] == sorted(active_experiment_names, reverse=True)[:3]
 
 
+def test_search_registered_models(tmp_path):
+    sqlite_uri = "sqlite:///{}".format(tmp_path.joinpath("test.db"))
+    mlflow.set_tracking_uri(sqlite_uri)
+
+    num_all_models = SEARCH_REGISTERED_MODEL_MAX_RESULTS_DEFAULT + 1
+    num_a_models = num_all_models // 4
+    num_b_models = num_all_models - num_a_models
+
+    a_model_names = [f"AModel_{i}" for i in range(num_a_models)]
+    b_model_names = [f"BModel_{i}" for i in range(num_b_models)]
+    model_names = b_model_names + a_model_names
+
+    tag_values = ["x", "x", "y"]
+    for (tag, model_name) in zip_longest(tag_values, model_names):
+        MlflowClient().create_registered_model(model_name, tags={"tag": tag} if tag else None)
+
+    # max_results is unspecified
+    models = mlflow.search_registered_models()
+    assert len(models) == num_all_models
+
+    # max_results is larger than the number of models in the database
+    models = mlflow.search_registered_models(max_results=num_all_models + 1)
+    assert len(models) == num_all_models
+
+    # max_results is equal to the number of models in the database
+    models = mlflow.search_registered_models(max_results=num_all_models)
+    assert len(models) == num_all_models
+    # max_results is smaller than the number of models in the database
+    models = mlflow.search_registered_models(max_results=num_all_models - 1)
+    assert len(models) == num_all_models - 1
+
+    # Filter by name
+    models = mlflow.search_registered_models(filter_string="name = 'AModel_1'")
+    assert [m.name for m in models] == ["AModel_1"]
+    models = mlflow.search_registered_models(filter_string="name ILIKE 'bmodel_%'")
+    assert len(models) == num_b_models
+
+    # Filter by tags
+    models = mlflow.search_registered_models(filter_string="tags.tag = 'x'")
+    assert [m.name for m in models] == model_names[:2]
+    models = mlflow.search_registered_models(filter_string="tags.tag = 'y'")
+    assert [m.name for m in models] == [model_names[2]]
+
+    # Order by name
+    models = mlflow.search_registered_models(order_by=["name DESC"], max_results=3)
+    assert [m.name for m in models] == sorted(model_names, reverse=True)[:3]
+
+
 @pytest.fixture
 def empty_active_run_stack():
     with mock.patch("mlflow.tracking.fluent._active_run_stack", []):