feat: support >1k series in fetch_metric_buckets (#107)

* feat: support >1k series in fetch_metric_buckets

* Un-skip tests for nan/inf

* Comment-out unused fields in TimeseriesBucket to save memory

* Remove TODO

* Fix fuzzy test after comment-out

* Restore constant

* More readable _update_range

Co-authored-by: Piotr Gabryjeluk <[email protected]>

* Fixes

* Fix min/max in update_range

---------

Co-authored-by: Piotr Gabryjeluk <[email protected]>

Author: pitercl

src/neptune_query/internal/composition/fetch_metric_buckets.py

@@ -18,9 +18,11 @@
     Generator,
     Literal,
     Optional,
+    Protocol,
     Union,
 )
 
+import numpy as np
 import pandas as pd
 from neptune_api.client import AuthenticatedClient
 
@@ -48,13 +50,23 @@
     util,
 )
 from ..retrieval.attribute_values import AttributeValue
-from ..retrieval.metric_buckets import TimeseriesBucket
+from ..retrieval.metric_buckets import (
+    MAX_SERIES_PER_REQUEST,
+    TimeseriesBucket,
+)
 from ..retrieval.search import ContainerType
 from .attribute_components import fetch_attribute_values_by_filter_split
 
 __all__ = ("fetch_metric_buckets",)
 
 
+class _FetchInChunksProtocol(Protocol):
+    def __call__(
+        self, x_range: Optional[tuple[float, float]], bucket_limit: int
+    ) -> dict[RunAttributeDefinition, list[TimeseriesBucket]]:
+        ...
+
+
 def fetch_metric_buckets(
     *,
     project_identifier: identifiers.ProjectIdentifier,
@@ -152,25 +164,102 @@ def go_fetch_sys_attrs() -> Generator[list[identifiers.SysId], None, None]:
         ),
     )
 
-    results: Generator[util.Page[AttributeValue], None, None] = concurrency.gather_results(output)
+    attribute_value_pages: Generator[util.Page[AttributeValue], None, None] = concurrency.gather_results(output)
 
     run_attribute_definitions = []
-    for page in results:
+    for page in attribute_value_pages:
         for value in page.items:
             run_attribute_definition = RunAttributeDefinition(
                 run_identifier=value.run_identifier, attribute_definition=value.attribute_definition
             )
             run_attribute_definitions.append(run_attribute_definition)
 
-    buckets_data = metric_buckets.fetch_time_series_buckets(
+    if not run_attribute_definitions:
+        return {}, sys_id_label_mapping
+
+    fetch_in_chunks: _FetchInChunksProtocol = lambda x_range, bucket_limit: _fetch_in_chunks(
         client=client,
-        x=x,
         run_attribute_definitions=run_attribute_definitions,
-        lineage_to_the_root=lineage_to_the_root,
+        x=x,
         include_point_previews=include_point_previews,
-        limit=limit,
         container_type=container_type,
-        x_range=None,
+        x_range=x_range,
+        lineage_to_the_root=lineage_to_the_root,
+        limit=bucket_limit,
+        executor=executor,
+    )
+
+    # if len(run_attribute_definitions) <= MAX_SERIES_PER_REQUEST:
+    #     return fetch_in_chunks(x_range=None, bucket_limit=limit), sys_id_label_mapping
+
+    global_x_range = _compute_global_x_range(fetch_in_chunks=fetch_in_chunks)
+    if global_x_range is None:
+        # No finite points / bucket bounds found
+        return {}, sys_id_label_mapping
+
+    return fetch_in_chunks(x_range=global_x_range, bucket_limit=limit), sys_id_label_mapping
+
+
+def _fetch_in_chunks(
+    client: AuthenticatedClient,
+    run_attribute_definitions: list[RunAttributeDefinition],
+    x: Literal["step"],
+    include_point_previews: bool,
+    container_type: ContainerType,
+    x_range: Optional[tuple[float, float]],
+    lineage_to_the_root: bool,
+    limit: int,
+    executor: Executor,
+) -> dict[RunAttributeDefinition, list[TimeseriesBucket]]:
+    chunks = (
+        run_attribute_definitions[offset : offset + MAX_SERIES_PER_REQUEST]
+        for offset in range(0, len(run_attribute_definitions), MAX_SERIES_PER_REQUEST)
     )
 
-    return buckets_data, sys_id_label_mapping
+    output = concurrency.generate_concurrently(
+        items=chunks,
+        executor=executor,
+        downstream=lambda chunk: concurrency.return_value(
+            metric_buckets.fetch_time_series_buckets(
+                client=client,
+                x=x,
+                run_attribute_definitions=chunk,
+                lineage_to_the_root=lineage_to_the_root,
+                include_point_previews=include_point_previews,
+                limit=limit,
+                container_type=container_type,
+                x_range=x_range,
+            )
+        ),
+    )
+
+    merged: dict[RunAttributeDefinition, list[TimeseriesBucket]] = {}
+    for chunk_data in concurrency.gather_results(output):  # type: dict[RunAttributeDefinition, list[TimeseriesBucket]]
+        merged.update(chunk_data)
+    return merged
+
+
+def _compute_global_x_range(fetch_in_chunks: _FetchInChunksProtocol) -> Optional[tuple[float, float]]:
+    x_range: tuple[Optional[float], Optional[float]] = (None, None)
+    # We only need the minimal number of buckets to determine min/max x
+    for buckets in fetch_in_chunks(x_range=None, bucket_limit=2).values():
+        for bucket in buckets:
+            x_range = _update_range(x_range, bucket)
+
+    if x_range[0] is None or x_range[1] is None:
+        return None
+    return x_range[0], x_range[1]
+
+
+def _update_range(
+    current_range: tuple[Optional[float], Optional[float]], bucket: TimeseriesBucket
+) -> tuple[Optional[float], Optional[float],]:
+    # We're including from_x and to_x because some buckets might hold only non-finite points,
+    # in which case first_x and last_x are None.
+    candidates = [bucket.first_x, bucket.last_x, bucket.from_x, bucket.to_x] + list(current_range)
+    finite_candidates = [x for x in candidates if x is not None and np.isfinite(x)]
+
+    if len(finite_candidates):
+        return min(finite_candidates), max(finite_candidates)
+    else:
+        return None, None

src/neptune_query/internal/retrieval/metric_buckets.py

@@ -44,6 +44,9 @@
 logger = get_logger()
 
 
+MAX_SERIES_PER_REQUEST = 1000
+
+
 @dataclass(frozen=True)
 class TimeseriesBucket:
     index: int
@@ -55,13 +58,13 @@ class TimeseriesBucket:
     last_y: Optional[float]
 
     # statistics:
-    y_min: Optional[float]
-    y_max: Optional[float]
-    finite_point_count: int
-    nan_count: int
-    positive_inf_count: int
-    negative_inf_count: int
-    finite_points_sum: Optional[float]
+    # y_min: Optional[float]
+    # y_max: Optional[float]
+    # finite_point_count: int
+    # nan_count: int
+    # positive_inf_count: int
+    # negative_inf_count: int
+    # finite_points_sum: Optional[float]
 
 
 # Build once at module import
@@ -112,8 +115,8 @@ def fetch_time_series_buckets(
     expressions = {}
     request_id_to_request_mapping = {}
     for num, run_attribute_definition in enumerate(run_attribute_definitions):
-        if num >= 1000:
-            raise ValueError("Cannot fetch more than 1000 time series at once")
+        if num >= MAX_SERIES_PER_REQUEST:
+            raise ValueError(f"Cannot fetch more than {MAX_SERIES_PER_REQUEST} time series at once")
 
         request_id = int_to_uuid(num)
         request_id_to_request_mapping[request_id] = run_attribute_definition
@@ -163,10 +166,10 @@ def fetch_time_series_buckets(
     for entry in result_object.entries:
         request = request_id_to_request_mapping.get(entry.requestId, None)
         if request is None:
-            raise RuntimeError(f"Received unknown requestId from the server: {request_id}")
+            raise RuntimeError(f"Received unknown requestId from the server: {entry.requestId}")
 
         if request in out:
-            raise RuntimeError(f"Received duplicate requestId from the server: {request_id}")
+            raise RuntimeError(f"Received duplicate requestId from the server: {entry.requestId}")
 
         out[request] = [
             TimeseriesBucket(
@@ -177,19 +180,19 @@ def fetch_time_series_buckets(
                 first_y=bucket.first.y if bucket.HasField("first") else None,
                 last_x=bucket.last.x if bucket.HasField("last") else None,
                 last_y=bucket.last.y if bucket.HasField("last") else None,
-                y_min=bucket.localMin if bucket.HasField("localMin") else None,
-                y_max=bucket.localMax if bucket.HasField("localMax") else None,
-                finite_point_count=bucket.finitePointCount,
-                nan_count=bucket.nanCount,
-                positive_inf_count=bucket.positiveInfCount,
-                negative_inf_count=bucket.negativeInfCount,
-                finite_points_sum=bucket.localSum if bucket.HasField("localSum") else None,
+                # y_min=bucket.localMin if bucket.HasField("localMin") else None,
+                # y_max=bucket.localMax if bucket.HasField("localMax") else None,
+                # finite_point_count=bucket.finitePointCount,
+                # nan_count=bucket.nanCount,
+                # positive_inf_count=bucket.positiveInfCount,
+                # negative_inf_count=bucket.negativeInfCount,
+                # finite_points_sum=bucket.localSum if bucket.HasField("localSum") else None,
             )
             for bucket in entry.bucket
         ]
 
     for request in run_attribute_definitions:
         if request not in out:
-            raise RuntimeError("Didn't get data for all the requests from the server. " f"Missing request {request_id}")
+            raise RuntimeError("Didn't get data for all the requests from the server. " f"Missing request {request}")
 
     return out

tests/e2e/metric_buckets.py

@@ -71,13 +71,13 @@ def aggregate_metric_buckets(
             first_y=ys[0] if ys else float("nan"),
             last_x=xs[-1] if xs else float("nan"),
             last_y=ys[-1] if ys else float("nan"),
-            y_min=float(np.min(ys)) if ys else float("nan"),
-            y_max=float(np.max(ys)) if ys else float("nan"),
-            finite_point_count=len(ys),
-            nan_count=nan_count,
-            positive_inf_count=positive_inf_count,
-            negative_inf_count=negative_inf_count,
-            finite_points_sum=float(np.sum(ys)) if ys else 0.0,
+            # y_min=float(np.min(ys)) if ys else float("nan"),
+            # y_max=float(np.max(ys)) if ys else float("nan"),
+            # finite_point_count=len(ys),
+            # nan_count=nan_count,
+            # positive_inf_count=positive_inf_count,
+            # negative_inf_count=negative_inf_count,
+            # finite_points_sum=float(np.sum(ys)) if ys else 0.0,
         )
         buckets.append(bucket)
     return buckets

tests/e2e/v1/test_experiments.py

@@ -321,7 +321,6 @@ def test__fetch_experiments_table_with_attributes_regex_filter_for_metrics(
     assert df[expected.columns].columns.equals(expected.columns)
 
 
-@pytest.mark.skip(reason="Skipped until inf/nan handling is enabled in the backend")
 def test__fetch_experiments_table_nan_inf(new_project_id):
     df = fetch_experiments_table(
         project=new_project_id,

tests/e2e/v1/test_fetch_metric_buckets.py

@@ -1,3 +1,4 @@
+import threading
 from typing import (
     Iterable,
     Literal,
@@ -21,6 +22,7 @@
     SysId,
 )
 from neptune_query.internal.output_format import create_metric_buckets_dataframe
+from neptune_query.internal.retrieval import metric_buckets
 from neptune_query.internal.retrieval.metric_buckets import TimeseriesBucket
 from tests.e2e.data import (
     NUMBER_OF_STEPS,
@@ -296,6 +298,85 @@ def test__fetch_metric_buckets__handles_misaligned_steps_in_metrics(
     pd.testing.assert_frame_equal(result_df, expected_df)
 
 
+@pytest.mark.parametrize(
+    "attribute_filter, expected_attributes",
+    [
+        (
+            AttributeFilter(name=r"series-.*", type=["float_series"]),
+            [
+                "series-containing-inf",
+                "series-containing-nan",
+                "series-ending-with-inf",
+                "series-ending-with-nan",
+            ],
+        ),
+        (
+            r"series-ending-.*",
+            ["series-ending-with-inf", "series-ending-with-nan"],
+        ),
+    ],
+)
+def test__fetch_metric_buckets__over_1k_series(
+    new_project_id,
+    monkeypatch,
+    attribute_filter,
+    expected_attributes,
+):
+    """
+    This test verifies that when fetching metric buckets for a run with over 1000 series,
+    the function correctly splits the requests into multiple chunks to avoid exceeding the limit.
+
+    It does so by monkeypatching the actual limit to 1 and capturing the calls made to fetch_time_series_buckets.
+    """
+    original_fetch = metric_buckets.fetch_time_series_buckets
+    call_chunks: list[list[RunAttributeDefinition]] = []
+    lock = threading.Lock()
+
+    def capture_fetch(*args, **kwargs):
+        with lock:
+            call_chunks.append(kwargs["run_attribute_definitions"])
+        return original_fetch(*args, **kwargs)
+
+    monkeypatch.setattr(metric_buckets, "fetch_time_series_buckets", capture_fetch)
+
+    forced_limit = 1
+    monkeypatch.setattr("neptune_query.internal.retrieval.metric_buckets.MAX_SERIES_PER_REQUEST", forced_limit)
+    monkeypatch.setattr("neptune_query.internal.composition.fetch_metric_buckets.MAX_SERIES_PER_REQUEST", forced_limit)
+
+    experiment_name = EXP_NAME_INF_NAN_RUN
+    run_id = RUN_ID_INF_NAN_RUN
+
+    result_df = fetch_metric_buckets(
+        project=new_project_id,
+        experiments=[experiment_name],
+        x="step",
+        y=attribute_filter,
+        limit=5,
+        include_point_previews=False,
+        lineage_to_the_root=True,
+    )
+
+    expected_data = {
+        experiment_name: {
+            attribute_name: RUN_BY_ID[run_id].metrics_values(attribute_name) for attribute_name in expected_attributes
+        }
+    }
+    expected_df = _create_expected_data_metric_buckets_dataframe(
+        data=expected_data,
+        project_identifier=new_project_id,
+        x="step",
+        limit=5,
+        include_point_previews=False,
+    )
+
+    pd.testing.assert_frame_equal(result_df, expected_df)
+
+    assert len(call_chunks) > 1
+    total_series = sum(len(chunk) for chunk in call_chunks)
+    assert total_series > forced_limit
+    assert all(len(chunk) <= forced_limit for chunk in call_chunks)
+
+
 @pytest.mark.parametrize(
     "arg_experiments,run_id,y",
     [
@@ -317,7 +398,6 @@ def test__fetch_metric_buckets__handles_misaligned_steps_in_metrics(
     "include_point_previews",
     [True],
 )
-@pytest.mark.skip(reason="Skipped until inf/nan handling is enabled in the backend")
 def test__fetch_metric_buckets__inf_nan(
     new_project_id,
     arg_experiments,

tests/e2e/v1/test_fetch_metrics.py

@@ -405,7 +405,6 @@ def test__fetch_metrics__lineage(new_project_id, lineage_to_the_root, expected_v
         ("series-containing-nan", RUN_BY_ID[RUN_ID_INF_NAN_RUN].metrics_values("series-containing-nan")),
     ],
 )
-@pytest.mark.skip(reason="Skipped until inf/nan handling is enabled in the backend")
 def test__fetch_metrics_nan_inf(new_project_id, series_name, expected_values):
     df = fetch_metrics(
         project=new_project_id,