Improve memory efficiency of metrics retrieval

Author: pitercl

src/neptune_query/internal/composition/fetch_metrics.py

@@ -46,7 +46,7 @@
     split,
 )
 from ..retrieval.metrics import (
-    FloatPointValue,
+    MetricValues,
     fetch_multiple_series_values,
 )
 from ..retrieval.search import ContainerType
@@ -96,6 +96,7 @@ def fetch_metrics(
             project_identifier=project_identifier,
             step_range=step_range,
             lineage_to_the_root=lineage_to_the_root,
+            include_timestamp=include_time is not None,
             include_point_previews=include_point_previews,
             tail_limit=tail_limit,
             executor=executor,
@@ -124,10 +125,11 @@ def _fetch_metrics(
     fetch_attribute_definitions_executor: Executor,
     step_range: tuple[Optional[float], Optional[float]],
     lineage_to_the_root: bool,
+    include_timestamp: bool,
     include_point_previews: bool,
     tail_limit: Optional[int],
     container_type: ContainerType,
-) -> tuple[dict[identifiers.RunAttributeDefinition, list[FloatPointValue]], dict[identifiers.SysId, str]]:
+) -> tuple[dict[identifiers.RunAttributeDefinition, MetricValues], dict[identifiers.SysId, str]]:
     sys_id_label_mapping: dict[identifiers.SysId, str] = {}
 
     def go_fetch_sys_attrs() -> Generator[list[identifiers.SysId], None, None]:
@@ -170,6 +172,7 @@ def go_fetch_sys_attrs() -> Generator[list[identifiers.SysId], None, None]:
                         client=client,
                         run_attribute_definitions=run_attribute_definitions_split,
                         include_inherited=lineage_to_the_root,
+                        include_timestamp=include_timestamp,
                         include_preview=include_point_previews,
                         container_type=container_type,
                         step_range=step_range,
@@ -180,13 +183,12 @@ def go_fetch_sys_attrs() -> Generator[list[identifiers.SysId], None, None]:
         ),
     )
 
-    results: Generator[
-        dict[identifiers.RunAttributeDefinition, list[FloatPointValue]], None, None
-    ] = concurrency.gather_results(output)
+    results: Generator[dict[identifiers.RunAttributeDefinition, MetricValues], None, None] = concurrency.gather_results(
+        output
+    )
 
-    metrics_data: dict[identifiers.RunAttributeDefinition, list[FloatPointValue]] = {}
-    for result in results:
-        for run_attribute_definition, metric_points in result.items():
-            metrics_data.setdefault(run_attribute_definition, []).extend(metric_points)
+    metrics_data: dict[identifiers.RunAttributeDefinition, MetricValues] = {
+        definition: metric_values for result in results for definition, metric_values in result.items()
+    }
 
     return metrics_data, sys_id_label_mapping

src/neptune_query/internal/output_format.py

@@ -31,7 +31,6 @@
 from . import identifiers
 from .retrieval import (
     metric_buckets,
-    metrics,
     series,
 )
 from .retrieval.attribute_types import (
@@ -40,13 +39,7 @@
     Histogram,
 )
 from .retrieval.attribute_values import AttributeValue
-from .retrieval.metrics import (
-    IsPreviewIndex,
-    PreviewCompletionIndex,
-    StepIndex,
-    TimestampIndex,
-    ValueIndex,
-)
+from .retrieval.metrics import MetricValues
 from .retrieval.search import ContainerType
 from .util import _validate_allowed_value
 
@@ -142,7 +135,7 @@ def transform_column_names(df: pd.DataFrame) -> pd.DataFrame:
 
 
 def create_metrics_dataframe(
-    metrics_data: dict[identifiers.RunAttributeDefinition, list[metrics.FloatPointValue]],
+    metrics_data: dict[identifiers.RunAttributeDefinition, MetricValues],
     sys_id_label_mapping: dict[identifiers.SysId, str],
     *,
     type_suffix_in_column_names: bool,
@@ -176,15 +169,13 @@ def path_display_name(attr_def: identifiers.RunAttributeDefinition) -> str:
     paths_with_data: set[str] = set()
 
     # Collect which (experiment, path) pairs have data and the set of observed steps per run.
-    for definition, points in metrics_data.items():
-        if not points:
+    for definition, metric_values in metrics_data.items():
+        if metric_values.length == 0:
             continue
 
         paths_with_data.add(path_display_name(definition))
-
         step_set = run_to_observed_steps.setdefault(sys_id_label_mapping[definition.run_identifier.sys_id], set())
-        for point in points:
-            step_set.add(point[StepIndex])
+        step_set.update(metric_values.steps)
 
     index_data = IndexData.from_observed_steps(
         observed_steps=run_to_observed_steps,
@@ -205,21 +196,21 @@ def path_display_name(attr_def: identifiers.RunAttributeDefinition) -> str:
     )
 
     # Write every metric point directly into the pre-allocated buffers.
-    for definition, points in metrics_data.items():
-        if not points:
+    for definition, metric_values in metrics_data.items():
+        if metric_values.length == 0:
             continue
 
         step_to_row_index: dict[float, int] = index_data.lookup_rows(sys_id=definition.run_identifier.sys_id)
+        rows = np.array([step_to_row_index[step] for step in metric_values.steps], dtype=np.uint)
+
         buffer: PathBuffer = path_buffers[path_display_name(definition)]
-        for point in points:
-            row_idx: int = step_to_row_index[point[StepIndex]]
-            buffer.value[row_idx] = point[ValueIndex]
-            if buffer.absolute_time is not None:
-                buffer.absolute_time[row_idx] = point[TimestampIndex]
-            if buffer.is_preview is not None:
-                buffer.is_preview[row_idx] = point[IsPreviewIndex]
-            if buffer.preview_completion is not None:
-                buffer.preview_completion[row_idx] = point[PreviewCompletionIndex]
+        buffer.value[rows] = metric_values.values
+        if buffer.absolute_time is not None:
+            buffer.absolute_time[rows] = metric_values.timestamps
+        if buffer.is_preview is not None:
+            buffer.is_preview[rows] = metric_values.is_preview
+        if buffer.preview_completion is not None:
+            buffer.preview_completion[rows] = metric_values.completion_ratio
 
     return _assemble_wide_dataframe(
         index_data=index_data,

src/neptune_query/internal/retrieval/metrics.py

@@ -14,12 +14,14 @@
 # limitations under the License.
 
 import functools as ft
+from dataclasses import dataclass
 from typing import (
     Any,
     Optional,
     Union,
 )
 
+import numpy as np
 from neptune_api.api.retrieval import get_multiple_float_series_values_proto
 from neptune_api.client import AuthenticatedClient
 from neptune_api.models import FloatTimeSeriesValuesRequest
@@ -37,28 +39,62 @@
 
 logger = get_logger()
 
-# Tuples are used here to enhance performance
-FloatPointValue = tuple[float, float, float, bool, float]
-(
-    TimestampIndex,
-    StepIndex,
-    ValueIndex,
-    IsPreviewIndex,
-    PreviewCompletionIndex,
-) = range(5)
-
 TOTAL_POINT_LIMIT: int = 1_000_000
 
 
+@dataclass(frozen=True, slots=True)
+class MetricValues:
+    steps: np.ndarray
+    values: np.ndarray
+    timestamps: Optional[np.ndarray]
+    is_preview: Optional[np.ndarray]
+    completion_ratio: Optional[np.ndarray]
+
+    @classmethod
+    def allocate(cls, size: int, include_timestamp: bool, include_preview: bool) -> "MetricValues":
+        return cls(
+            steps=np.empty(size, dtype=np.float64),
+            values=np.empty(size, dtype=np.float64),
+            timestamps=np.empty(size, dtype=np.float64) if include_timestamp else None,
+            is_preview=np.empty(size, dtype=bool) if include_preview else None,
+            completion_ratio=np.empty(size, dtype=np.float64) if include_preview else None,
+        )
+
+    @classmethod
+    def concatenate(cls, metrics_list: list["MetricValues"]) -> "MetricValues":
+        return cls(
+            steps=np.concatenate([m.steps for m in metrics_list], axis=0),
+            values=np.concatenate([m.values for m in metrics_list], axis=0),
+            timestamps=np.concatenate([m.timestamps for m in metrics_list], axis=0)
+            if metrics_list[0].timestamps is not None
+            else None,
+            is_preview=np.concatenate([m.is_preview for m in metrics_list], axis=0)
+            if metrics_list[0].is_preview is not None
+            else None,
+            completion_ratio=np.concatenate([m.completion_ratio for m in metrics_list], axis=0)
+            if metrics_list[0].completion_ratio is not None
+            else None,
+        )
+
+    @property
+    def length(self) -> int:
+        return len(self.steps)
+
+    @classmethod
+    def length_sum(cls, metrics_list: list["MetricValues"]) -> int:
+        return sum(m.length for m in metrics_list)
+
+
 def fetch_multiple_series_values(
     client: AuthenticatedClient,
     run_attribute_definitions: list[identifiers.RunAttributeDefinition],
     include_inherited: bool,
     container_type: ContainerType,
+    include_timestamp: bool,
     include_preview: bool,
     step_range: tuple[Union[float, None], Union[float, None]] = (None, None),
     tail_limit: Optional[int] = None,
-) -> dict[identifiers.RunAttributeDefinition, list[FloatPointValue]]:
+) -> dict[identifiers.RunAttributeDefinition, MetricValues]:
     if not run_attribute_definitions:
         return {}
 
@@ -93,25 +129,37 @@ def fetch_multiple_series_values(
         "order": "ascending" if not tail_limit else "descending",
     }
 
-    results: dict[identifiers.RunAttributeDefinition, list[FloatPointValue]] = {
-        run_attribute: [] for run_attribute in run_attribute_definitions
-    }
+    paged_results: dict[identifiers.RunAttributeDefinition, list[MetricValues]] = {}
 
     for page_result in util.fetch_pages(
         client=client,
         fetch_page=_fetch_metrics_page,
-        process_page=ft.partial(_process_metrics_page, request_id_to_attribute=request_id_to_attribute),
+        process_page=ft.partial(
+            _process_metrics_page,
+            request_id_to_attribute=request_id_to_attribute,
+            include_timestamp=include_timestamp,
+            include_preview=include_preview,
+            reverse_order=tail_limit is not None,
+        ),
         make_new_page_params=ft.partial(
             _make_new_metrics_page_params,
             request_id_to_attribute=request_id_to_attribute,
             tail_limit=tail_limit,
-            partial_results=results,
+            partial_results=paged_results,
         ),
         params=params,
     ):
-        for attribute, values in page_result.items:
-            sorted_values = values if tail_limit else reversed(values)
-            results[attribute].extend(sorted_values)
+        for definition, metric_values in page_result.items:
+            paged_results.setdefault(definition, []).append(metric_values)
+
+    results: dict[identifiers.RunAttributeDefinition, MetricValues] = {}
+    for definition, paged_metric_values in paged_results.items():
+        if len(paged_metric_values) > 1:
+            results[definition] = MetricValues.concatenate(paged_metric_values)
+        elif len(paged_metric_values) == 1:
+            results[definition] = paged_metric_values[0]
+        else:
+            pass
 
     return results
 
@@ -138,20 +186,32 @@ def _fetch_metrics_page(
 def _process_metrics_page(
     data: ProtoFloatSeriesValuesResponseDTO,
     request_id_to_attribute: dict[str, identifiers.RunAttributeDefinition],
-) -> util.Page[tuple[identifiers.RunAttributeDefinition, list[FloatPointValue]]]:
+    include_timestamp: bool,
+    include_preview: bool,
+    reverse_order: bool,
+) -> util.Page[tuple[identifiers.RunAttributeDefinition, MetricValues]]:
     result = {}
     for series in data.series:
-        run_attribute = request_id_to_attribute[series.requestId]
-        result[run_attribute] = [
-            (
-                point.timestamp_millis,
-                point.step,
-                point.value,
-                point.is_preview,
-                point.completion_ratio,
-            )
-            for point in series.series.values
-        ]
+        metric_values = MetricValues.allocate(
+            size=len(series.series.values), include_timestamp=include_timestamp, include_preview=include_preview
+        )
+
+        for i, point in enumerate(series.series.values):
+            idx = metric_values.length - 1 - i if reverse_order else i
+
+            metric_values.steps[idx] = point.step
+            metric_values.values[idx] = point.value
+            if include_timestamp:
+                assert metric_values.timestamps
+                metric_values.timestamps[idx] = point.timestamp_millis
+            if include_preview:
+                assert metric_values.is_preview
+                assert metric_values.completion_ratio
+                metric_values.is_preview[idx] = point.is_preview
+                metric_values.completion_ratio[idx] = point.completion_ratio
+        definition = request_id_to_attribute[series.requestId]
+        result[definition] = metric_values
+
     return util.Page(items=list(result.items()))
 
 
@@ -160,7 +220,7 @@ def _make_new_metrics_page_params(
     data: Optional[ProtoFloatSeriesValuesResponseDTO],
     request_id_to_attribute: dict[str, identifiers.RunAttributeDefinition],
     tail_limit: Optional[int],
-    partial_results: dict[identifiers.RunAttributeDefinition, list[FloatPointValue]],
+    partial_results: dict[identifiers.RunAttributeDefinition, list[MetricValues]],
 ) -> Optional[dict[str, Any]]:
     if data is None:  # no past data, we are fetching the first page
         for request in params["requests"]:
@@ -181,7 +241,9 @@ def _make_new_metrics_page_params(
         is_page_full = value_size == prev_per_series_points_limit
 
         attribute = request_id_to_attribute[request_id]
-        need_more_points = len(partial_results[attribute]) < tail_limit if tail_limit is not None else True
+        need_more_points = (
+            MetricValues.length_sum(partial_results[attribute]) < tail_limit if tail_limit is not None else True
+        )
 
         if is_page_full and need_more_points:
             new_request_after_steps[request_id] = series.series.values[-1].step
@@ -201,7 +263,8 @@ def _make_new_metrics_page_params(
     per_series_points_limit = max(1, TOTAL_POINT_LIMIT // len(params["requests"]))
     if tail_limit is not None:
         already_fetched = next(
-            len(partial_results[request_id_to_attribute[request_id]]) for request_id in new_request_after_steps.keys()
+            MetricValues.length_sum(partial_results[request_id_to_attribute[request_id]])
+            for request_id in new_request_after_steps.keys()
         )  # assumes the results for all unfinished series have the same length
         per_series_points_limit = min(per_series_points_limit, tail_limit - already_fetched)
     params["perSeriesPointsLimit"] = per_series_points_limit