Search Logged Models Support Datasets Filter for File Store and Sqlalchemy Store (mlflow#16262)

Signed-off-by: Raymond Zhou <[email protected]>

Author: raymondzhou-db

mlflow/server/handlers.py

@@ -2827,6 +2827,7 @@ def _search_logged_models():
                 _assert_required,
             ],
             "filter": [_assert_string],
+            "datasets": [_assert_array],
             "max_results": [_assert_intlike],
             "order_by": [_assert_array],
             "page_token": [_assert_string],
@@ -2837,6 +2838,17 @@ def _search_logged_models():
         # to avoid serialization issues
         experiment_ids=list(request_message.experiment_ids),
         filter_string=request_message.filter or None,
+        datasets=(
+            [
+                {
+                    "dataset_name": d.dataset_name,
+                    "dataset_digest": d.dataset_digest or None,
+                }
+                for d in request_message.datasets
+            ]
+            if request_message.datasets
+            else None
+        ),
         max_results=request_message.max_results or None,
         order_by=(
             [

mlflow/store/tracking/file_store.py

@@ -8,7 +8,7 @@
 import uuid
 from collections import defaultdict
 from dataclasses import dataclass
-from typing import Any, NamedTuple, Optional
+from typing import Any, NamedTuple, Optional, TypedDict
 
 from mlflow.entities import (
     Dataset,
@@ -165,6 +165,15 @@ def _read_persisted_run_info_dict(run_info_dict):
     return RunInfo.from_dictionary(dict_copy)
 
 
+class DatasetFilter(TypedDict, total=False):
+    """
+    Dataset filter used for search_logged_models.
+    """
+
+    dataset_name: str
+    dataset_digest: str
+
+
 class FileStore(AbstractStore):
     TRASH_FOLDER_NAME = ".trash"
     ARTIFACTS_FOLDER_NAME = "artifacts"
@@ -2285,7 +2294,7 @@ def search_logged_models(
         self,
         experiment_ids: list[str],
         filter_string: Optional[str] = None,
-        datasets: Optional[list[str]] = None,
+        datasets: Optional[list[DatasetFilter]] = None,
         max_results: Optional[int] = None,
         order_by: Optional[list[dict[str, Any]]] = None,
         page_token: Optional[str] = None,
@@ -2299,8 +2308,8 @@ def search_logged_models(
             datasets: List of dictionaries to specify datasets on which to apply metrics filters.
                 The following fields are supported:
 
-                name (str): Required. Name of the dataset.
-                digest (str): Optional. Digest of the dataset.
+                dataset_name (str): Required. Name of the dataset.
+                dataset_digest (str): Optional. Digest of the dataset.
             max_results: Maximum number of logged models desired. Default is 100.
             order_by: List of dictionaries to specify the ordering of the search results.
                 The following fields are supported:
@@ -2321,17 +2330,17 @@ def search_logged_models(
             A :py:class:`PagedList <mlflow.store.entities.PagedList>` of
             :py:class:`LoggedModel <mlflow.entities.LoggedModel>` objects.
         """
-        if datasets:
+        if datasets and not all(d.get("dataset_name") for d in datasets):
             raise MlflowException(
-                "Filtering by datasets is not currently supported by FileStore",
+                "`dataset_name` in the `datasets` clause must be specified.",
                 INVALID_PARAMETER_VALUE,
             )
         max_results = max_results or SEARCH_LOGGED_MODEL_MAX_RESULTS_DEFAULT
         all_models = []
         for experiment_id in experiment_ids:
             models = self._list_models(experiment_id)
             all_models.extend(models)
-        filtered = SearchLoggedModelsUtils.filter_logged_models(all_models, filter_string)
+        filtered = SearchLoggedModelsUtils.filter_logged_models(all_models, filter_string, datasets)
         sorted_logged_models = SearchLoggedModelsUtils.sort(filtered, order_by)
         logged_models, next_page_token = SearchLoggedModelsUtils.paginate(
             sorted_logged_models, page_token, max_results

mlflow/store/tracking/rest_store.py

@@ -1008,8 +1008,8 @@ def search_logged_models(
             datasets: List of dictionaries to specify datasets on which to apply metrics filters.
                 The following fields are supported:
 
-                name (str): Required. Name of the dataset.
-                digest (str): Optional. Digest of the dataset.
+                dataset_name (str): Required. Name of the dataset.
+                dataset_digest (str): Optional. Digest of the dataset.
             max_results: Maximum number of logged models desired.
             order_by: List of dictionaries to specify the ordering of the search results.
                 The following fields are supported:

mlflow/store/tracking/sqlalchemy_store.py

@@ -7,7 +7,7 @@
 import uuid
 from collections import defaultdict
 from functools import reduce
-from typing import Any, Optional
+from typing import Any, Optional, TypedDict
 
 import sqlalchemy
 import sqlalchemy.orm
@@ -126,6 +126,15 @@
 sqlalchemy.orm.configure_mappers()
 
 
+class DatasetFilter(TypedDict, total=False):
+    """
+    Dataset filter used for search_logged_models.
+    """
+
+    dataset_name: str
+    dataset_digest: str
+
+
 class SqlAlchemyStore(AbstractStore):
     """
     SQLAlchemy compliant backend store for tracking meta data for MLflow entities. MLflow
@@ -1968,31 +1977,46 @@ def _apply_order_by_search_logged_models(
 
         return models.order_by(*order_by_clauses)
 
-    def _apply_filter_string_search_logged_models(
+    def _apply_filter_string_datasets_search_logged_models(
         self,
         models: sqlalchemy.orm.Query,
         session: sqlalchemy.orm.Session,
         experiment_ids: list[str],
         filter_string: Optional[str],
+        datasets: Optional[list[dict[str, Any]]],
     ):
         from mlflow.utils.search_logged_model_utils import EntityType, parse_filter_string
 
         comparisons = parse_filter_string(filter_string)
         dialect = self._get_dialect()
         attr_filters: list[sqlalchemy.BinaryExpression] = []
         non_attr_filters: list[sqlalchemy.BinaryExpression] = []
+
+        dataset_filters = []
+        if datasets:
+            for dataset in datasets:
+                dataset_filter = SqlLoggedModelMetric.dataset_name == dataset["dataset_name"]
+                if "dataset_digest" in dataset:
+                    dataset_filter = dataset_filter & (
+                        SqlLoggedModelMetric.dataset_digest == dataset["dataset_digest"]
+                    )
+                dataset_filters.append(dataset_filter)
+
+        has_metric_filters = False
         for comp in comparisons:
             comp_func = SearchUtils.get_sql_comparison_func(comp.op, dialect)
             if comp.entity.type == EntityType.ATTRIBUTE:
                 attr_filters.append(comp_func(getattr(SqlLoggedModel, comp.entity.key), comp.value))
             elif comp.entity.type == EntityType.METRIC:
+                has_metric_filters = True
+                metric_filters = [
+                    SqlLoggedModelMetric.metric_name == comp.entity.key,
+                    comp_func(SqlLoggedModelMetric.metric_value, comp.value),
+                ]
+                if dataset_filters:
+                    metric_filters.append(sqlalchemy.or_(*dataset_filters))
                 non_attr_filters.append(
-                    session.query(SqlLoggedModelMetric)
-                    .filter(
-                        SqlLoggedModelMetric.metric_name == comp.entity.key,
-                        comp_func(SqlLoggedModelMetric.metric_value, comp.value),
-                    )
-                    .subquery()
+                    session.query(SqlLoggedModelMetric).filter(*metric_filters).subquery()
                 )
             elif comp.entity.type == EntityType.PARAM:
                 non_attr_filters.append(
@@ -2016,6 +2040,17 @@ def _apply_filter_string_search_logged_models(
         for f in non_attr_filters:
             models = models.join(f)
 
+        # If there are dataset filters but no metric filters,
+        # filter for models that have any metrics on the datasets
+        if dataset_filters and not has_metric_filters:
+            subquery = (
+                session.query(SqlLoggedModelMetric.model_id)
+                .filter(sqlalchemy.or_(*dataset_filters))
+                .distinct()
+                .subquery()
+            )
+            models = models.join(subquery)
+
         return models.filter(
             SqlLoggedModel.lifecycle_stage != LifecycleStage.DELETED,
             SqlLoggedModel.experiment_id.in_(experiment_ids),
@@ -2026,14 +2061,14 @@ def search_logged_models(
         self,
         experiment_ids: list[str],
         filter_string: Optional[str] = None,
-        datasets: Optional[list[DatasetInput]] = None,
+        datasets: Optional[list[DatasetFilter]] = None,
         max_results: Optional[int] = None,
         order_by: Optional[list[dict[str, Any]]] = None,
         page_token: Optional[str] = None,
     ) -> PagedList[LoggedModel]:
-        if datasets:
+        if datasets and not all(d.get("dataset_name") for d in datasets):
             raise MlflowException(
-                "Filtering by datasets is not currently supported by SqlAlchemyStore",
+                "`dataset_name` in the `datasets` clause must be specified.",
                 INVALID_PARAMETER_VALUE,
             )
         if page_token:
@@ -2046,8 +2081,8 @@ def search_logged_models(
         max_results = max_results or SEARCH_LOGGED_MODEL_MAX_RESULTS_DEFAULT
         with self.ManagedSessionMaker() as session:
             models = session.query(SqlLoggedModel)
-            models = self._apply_filter_string_search_logged_models(
-                models, session, experiment_ids, filter_string
+            models = self._apply_filter_string_datasets_search_logged_models(
+                models, session, experiment_ids, filter_string, datasets
             )
             models = self._apply_order_by_search_logged_models(models, session, order_by)
             models = models.offset(offset).limit(max_results + 1).all()

mlflow/utils/search_utils.py

@@ -20,7 +20,7 @@
 )
 from sqlparse.tokens import Token as TokenType
 
-from mlflow.entities import LoggedModel, RunInfo
+from mlflow.entities import LoggedModel, Metric, RunInfo
 from mlflow.entities.model_registry.model_version_stages import STAGE_DELETED_INTERNAL
 from mlflow.entities.model_registry.prompt_version import IS_PROMPT_TAG_KEY
 from mlflow.exceptions import MlflowException
@@ -593,6 +593,13 @@ def is_dataset(cls, key_type, comparator):
             return True
         return False
 
+    @classmethod
+    def _is_metric_on_dataset(cls, metric: Metric, dataset: dict[str, Any]) -> bool:
+        return metric.dataset_name == dataset.get("dataset_name") and (
+            dataset.get("dataset_digest") is None
+            or dataset.get("dataset_digest") == metric.dataset_digest
+        )
+
     @classmethod
     def _does_run_match_clause(cls, run, sed):
         key_type = sed.get("type")
@@ -1859,7 +1866,12 @@ class SearchLoggedModelsUtils(SearchUtils):
     VALID_ORDER_BY_ATTRIBUTE_KEYS = VALID_SEARCH_ATTRIBUTE_KEYS
 
     @classmethod
-    def _does_logged_model_match_clause(cls, model: LoggedModel, condition: dict[str, Any]):
+    def _does_logged_model_match_clause(
+        cls,
+        model: LoggedModel,
+        condition: dict[str, Any],
+        datasets: Optional[list[dict[str, Any]]] = None,
+    ):
         key_type = condition.get("type")
         key = condition.get("key")
         value = condition.get("value")
@@ -1869,6 +1881,12 @@ def _does_logged_model_match_clause(cls, model: LoggedModel, condition: dict[str
 
         if cls.is_metric(key_type, comparator):
             matching_metrics = [metric for metric in model.metrics if metric.key == key]
+            if datasets:
+                matching_metrics = [
+                    metric
+                    for metric in matching_metrics
+                    if any(cls._is_metric_on_dataset(metric, dataset) for dataset in datasets)
+                ]
             lhs = matching_metrics[0].value if matching_metrics else None
             value = float(value)
         elif cls.is_param(key_type, comparator):
@@ -1896,15 +1914,34 @@ def validate_list_supported(cls, key: str) -> None:
         """
 
     @classmethod
-    def filter_logged_models(cls, models: list[LoggedModel], filter_string: Optional[str] = None):
-        """Filters a set of runs based on a search filter string."""
-        if not filter_string:
+    def filter_logged_models(
+        cls,
+        models: list[LoggedModel],
+        filter_string: Optional[str] = None,
+        datasets: Optional[list[dict[str, Any]]] = None,
+    ):
+        """Filters a set of runs based on a search filter string and list of dataset filters."""
+        if not filter_string and not datasets:
             return models
 
         parsed = cls.parse_search_filter(filter_string)
 
+        # If there are dataset filters but no metric filters in the filter string,
+        # filter for models that have any metrics on the datasets
+        if datasets and not any(
+            cls.is_metric(s.get("type"), s.get("comparator").upper()) for s in parsed
+        ):
+
+            def model_has_metrics_on_datasets(model):
+                return any(
+                    any(cls._is_metric_on_dataset(metric, dataset) for dataset in datasets)
+                    for metric in model.metrics
+                )
+
+            models = [model for model in models if model_has_metrics_on_datasets(model)]
+
         def model_matches(model):
-            return all(cls._does_logged_model_match_clause(model, s) for s in parsed)
+            return all(cls._does_logged_model_match_clause(model, s, datasets) for s in parsed)
 
         return [model for model in models if model_matches(model)]