[GSoC] Add paginated, filtered run iteration to the core client

metaflow/client/core.py

@@ -4,10 +4,11 @@
 import os
 import tarfile
 from collections import namedtuple
+from collections.abc import Mapping
 from datetime import datetime
 from tempfile import TemporaryDirectory
 from io import BytesIO
-from itertools import chain
+from itertools import chain, islice
 from typing import (
     Any,
     Dict,
@@ -401,19 +402,7 @@ def __iter__(self) -> Iterator["MetaflowObject"]:
         unfiltered_children = unfiltered_children if unfiltered_children else []
         children = filter(
             lambda x: self._iter_filter(x),
-            (
-                _CLASSES[self._CHILD_CLASS](
-                    attempt=self._attempt,
-                    _object=obj,
-                    _parent=self,
-                    _metaflow=self._metaflow,
-                    _namespace_check=self._namespace_check,
-                    _current_namespace=(
-                        self._current_namespace if self._namespace_check else None
-                    ),
-                )
-                for obj in unfiltered_children
-            ),
+            (self._child_from_record(obj) for obj in unfiltered_children),
         )
 
         if children:
@@ -424,6 +413,23 @@ def __iter__(self) -> Iterator["MetaflowObject"]:
     def _iter_filter(self, x):
         return True
 
+    def _child_from_record(self, obj):
+        """Build a child object from a raw metadata record.
+
+        Shared by __iter__ and _iter_children so the materialized and streaming
+        listing paths cannot drift apart.
+        """
+        return _CLASSES[self._CHILD_CLASS](
+            attempt=self._attempt,
+            _object=obj,
+            _parent=self,
+            _metaflow=self._metaflow,
+            _namespace_check=self._namespace_check,
+            _current_namespace=(
+                self._current_namespace if self._namespace_check else None
+            ),
+        )
+
     def _filtered_children(self, *tags):
         """
         Returns an iterator over all children.
@@ -435,6 +441,28 @@ def _filtered_children(self, *tags):
             if all(tag in child.tags for tag in tags):
                 yield child
 
+    def _iter_children(self, query_filters=None, page_size=None, required_tags=()):
+        """Stream child records through the active metadata provider."""
+        namespace_filter = {}
+        if self._namespace_check and self._current_namespace:
+            namespace_filter = {"any_tags": self._current_namespace}
+
+        objects = self._metaflow.metadata.iter_objects(
+            self._NAME,
+            _CLASSES[self._CHILD_CLASS]._NAME,
+            namespace_filter,
+            self._attempt,
+            *self.path_components,
+            query_filters=query_filters,
+            page_size=page_size,
+        )
+        for obj in objects:
+            child = self._child_from_record(obj)
+            if self._iter_filter(child) and all(
+                tag in child.tags for tag in required_tags
+            ):
+                yield child
+
     def _ipython_key_completions_(self):
         """Returns available options for ipython auto-complete."""
         return [child.id for child in self._filtered_children()]
@@ -2585,7 +2613,12 @@ def latest_successful_run(self) -> Optional[Run]:
             if run.successful:
                 return run
 
-    def runs(self, *tags: str) -> Iterator[Run]:
+    def runs(
+        self,
+        *tags: str,
+        _filters: Optional[Dict[str, Any]] = None,
+        max_runs: Optional[int] = None,
+    ) -> Iterator[Run]:
         """
         Returns an iterator over all `Run`s of this flow.
 
@@ -2596,14 +2629,49 @@ def runs(self, *tags: str) -> Iterator[Run]:
         Parameters
         ----------
         tags : str
-            Tags to match.
+            Tags to match. Applied locally after listing, so this works with
+            local metadata as well as the metadata service.
+        _filters : dict, optional
+            Internal, unstable. Server-side run filters expressed in the metadata
+            service's ``field:operator`` grammar, for example
+            ``{"status:eq": "failed"}``. The grammar is specific to the metadata
+            service and has no meaning for local metadata, so it is deliberately
+            not part of the public interface -- the typed accessors built on top
+            of it are the supported way to filter. Requires a metadata service
+            with pagination and filtering support.
+        max_runs : int, optional
+            Maximum number of runs to yield, newest first. Supported by every
+            metadata provider.
 
         Yields
         ------
         Run
             `Run` objects in this flow.
         """
-        return self._filtered_children(*tags)
+        if _filters is not None and not isinstance(_filters, Mapping):
+            raise TypeError("_filters must be a mapping")
+        if max_runs is not None:
+            # bool is a subclass of int, so isinstance(True, int) is True; without
+            # the explicit bool check runs(max_runs=True) would silently mean 1.
+            if isinstance(max_runs, bool) or not isinstance(max_runs, int):
+                raise TypeError("max_runs must be an integer")
+            if max_runs < 0:
+                raise ValueError("max_runs must be non-negative")
+            if max_runs == 0:
+                return iter(())
+
+        if _filters is None and max_runs is None:
+            return self._filtered_children(*tags)
+
+        # Page size is a metadata-service transport detail, not a caller concern:
+        # this returns an iterator either way. It comes from METAFLOW_SERVICE_PAGE_SIZE.
+        runs = self._iter_children(
+            query_filters=dict(_filters) if _filters else None,
+            required_tags=tags,
+        )
+        if max_runs is None:
+            return runs
+        return islice(runs, max_runs)
 
     def __iter__(self) -> Iterator[Task]:
         """

metaflow/metadata_provider/metadata.py

@@ -6,7 +6,11 @@
 from itertools import chain
 
 from typing import List
-from metaflow.exception import MetaflowInternalError, MetaflowTaggingError
+from metaflow.exception import (
+    MetaflowException,
+    MetaflowInternalError,
+    MetaflowTaggingError,
+)
 from metaflow.tagging_util import validate_tag
 from metaflow.util import get_username, resolve_identity_as_tuple, is_stringish
 
@@ -348,6 +352,37 @@ def add_sticky_tags(self, tags=None, sys_tags=None):
         if sys_tags:
             self.sticky_sys_tags.update(sys_tags)
 
+    @classmethod
+    def _validate_object_query(cls, obj_type, sub_type):
+        """Reject nonsensical obj_type/sub_type combinations.
+
+        Shared by get_object and the streaming listing paths so every access,
+        materialized or paginated, enforces the same rules. Returns the
+        (type_order, sub_order) pair for callers that need it.
+        """
+        type_order = ObjectOrder.type_to_order(obj_type)
+        sub_order = ObjectOrder.type_to_order(sub_type)
+
+        if type_order is None:
+            raise MetaflowInternalError(msg="Cannot find type %s" % obj_type)
+        if type_order >= ObjectOrder.type_to_order("metadata"):
+            raise MetaflowInternalError(msg="Type %s is not allowed" % obj_type)
+
+        if sub_order is None:
+            raise MetaflowInternalError(msg="Cannot find subtype %s" % sub_type)
+
+        if type_order >= sub_order:
+            raise MetaflowInternalError(
+                msg="Subtype %s not allowed for %s" % (sub_type, obj_type)
+            )
+
+        # Metadata is always only at the task level
+        if sub_type == "metadata" and obj_type != "task":
+            raise MetaflowInternalError(
+                msg="Metadata can only be retrieved at the task level"
+            )
+        return type_order, sub_order
+
     @classmethod
     def get_object(cls, obj_type, sub_type, filters, attempt, *args):
         """Returns the requested object depending on obj_type and sub_type
@@ -397,27 +432,7 @@ def get_object(cls, obj_type, sub_type, filters, attempt, *args):
             object or list :
                 Depending on the call, the type of object return varies
         """
-        type_order = ObjectOrder.type_to_order(obj_type)
-        sub_order = ObjectOrder.type_to_order(sub_type)
-
-        if type_order is None:
-            raise MetaflowInternalError(msg="Cannot find type %s" % obj_type)
-        if type_order >= ObjectOrder.type_to_order("metadata"):
-            raise MetaflowInternalError(msg="Type %s is not allowed" % obj_type)
-
-        if sub_order is None:
-            raise MetaflowInternalError(msg="Cannot find subtype %s" % sub_type)
-
-        if type_order >= sub_order:
-            raise MetaflowInternalError(
-                msg="Subtype %s not allowed for %s" % (sub_type, obj_type)
-            )
-
-        # Metadata is always only at the task level
-        if sub_type == "metadata" and obj_type != "task":
-            raise MetaflowInternalError(
-                msg="Metadata can only be retrieved at the task level"
-            )
+        type_order, sub_order = cls._validate_object_query(obj_type, sub_type)
 
         if attempt is not None:
             try:
@@ -439,6 +454,45 @@ def get_object(cls, obj_type, sub_type, filters, attempt, *args):
             pre_filter, attempt_int
         )
 
+    @classmethod
+    def iter_objects(cls, obj_type, sub_type, filters, attempt, *args, **kwargs):
+        """Iterate over a collection returned by ``get_object``.
+
+        Providers can override this method to stream records without materializing the
+        complete collection. ``query_filters`` and ``page_size`` are optional provider
+        hints; the default implementation supports only unfiltered iteration.
+        """
+        query_filters = kwargs.pop("query_filters", None)
+        page_size = kwargs.pop("page_size", None)
+        if kwargs:
+            raise TypeError("Unexpected iterator options: %s" % ", ".join(kwargs))
+        if query_filters:
+            raise MetaflowException(
+                "Server-side metadata filters are not supported by the %s provider"
+                % cls.TYPE
+            )
+        if page_size is not None:
+            if isinstance(page_size, bool) or not isinstance(page_size, int):
+                raise TypeError("page_size must be an integer")
+            if page_size <= 0:
+                raise ValueError("page_size must be positive")
+
+        objects = cls.get_object(obj_type, sub_type, filters, attempt, *args)
+        if isinstance(objects, dict):
+            objects = [objects]
+        # Yield newest-first (descending ts_epoch). This mirrors the order the
+        # metadata service returns for paginated listings, so callers see a
+        # consistent order whether records came from the service or a legacy /
+        # local provider. ts_epoch may be missing on some records, so fall back
+        # to 0 to keep the sort total.
+        objects = sorted(
+            objects or [],
+            key=lambda obj: obj.get("ts_epoch") or 0,
+            reverse=True,
+        )
+        for obj in objects:
+            yield obj
+
     @classmethod
     def mutate_user_tags_for_run(
         cls, flow_id, run_id, tags_to_remove=None, tags_to_add=None

metaflow/metaflow_config.py

@@ -6,7 +6,11 @@
 
 from typing import Dict, List, Union, Tuple as TTuple
 from metaflow.exception import MetaflowException
-from metaflow.metaflow_config_funcs import from_conf, get_validate_choice_fn
+from metaflow.metaflow_config_funcs import (
+    from_conf,
+    get_validate_choice_fn,
+    get_validate_positive_int_fn,
+)
 
 # Recursive type alias for JSON, used by Runner API type mappings
 JSON = Union[Dict[str, "JSON"], List["JSON"], str, int, float, bool, None]
@@ -287,6 +291,14 @@
 ###
 SERVICE_URL = from_conf("SERVICE_URL")
 SERVICE_RETRY_COUNT = from_conf("SERVICE_RETRY_COUNT", 5)
+# Number of records the client asks the metadata service for per page when it
+# streams a collection (for example Flow.runs()). Page size is a transport
+# detail rather than a caller concern -- the client listing APIs hand back an
+# iterator either way -- so it is configured here instead of per call. The
+# service caps a page at 500 records; larger values are clamped client-side.
+SERVICE_PAGE_SIZE = int(
+    from_conf("SERVICE_PAGE_SIZE", 100, validate_fn=get_validate_positive_int_fn())
+)
 SERVICE_AUTH_KEY = from_conf("SERVICE_AUTH_KEY")
 SERVICE_HEADERS = from_conf("SERVICE_HEADERS", {})
 if SERVICE_AUTH_KEY is not None:

metaflow/metaflow_config_funcs.py

@@ -178,3 +178,23 @@ def _validate_choice(name, value):
             )
 
     return _validate_choice
+
+
+def get_validate_positive_int_fn():
+    """Returns a validate_fn for use with from_conf().
+    The validate_fn will check that a value is a positive integer. Values read
+    from the environment arrive as strings, so this accepts anything int() can
+    parse and rejects the rest with a readable message.
+    """
+
+    def _validate_positive_int(name, value):
+        try:
+            parsed = int(value)
+        except (TypeError, ValueError):
+            raise MetaflowException("%s must be an integer. Got '%s'." % (name, value))
+        if parsed <= 0:
+            raise MetaflowException(
+                "%s must be a positive integer. Got '%s'." % (name, value)
+            )
+
+    return _validate_positive_int