Fetch runs table without Bravado deserialization

CHANGELOG.md

@@ -1,3 +1,9 @@
+## [UNRELEASED] 1.8.6
+
+### Changes
+- Improved performance of `fetch_*_table()` methods up to 2x ([#1573])(https://github.com/neptune-ai/neptune-client/pull/1573)
+
+
 ## neptune 1.8.5
 
 ### Fixes

src/neptune/api/searching_entries.py

@@ -0,0 +1,108 @@
+#
+# Copyright (c) 2023, Neptune Labs Sp. z o.o.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+__all__ = ["get_single_page", "iter_over_pages", "to_leaderboard_entry"]
+
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Dict,
+    Generator,
+    Iterable,
+    List,
+    Optional,
+)
+
+from bravado.client import construct_request  # type: ignore
+from bravado.config import RequestConfig  # type: ignore
+
+from neptune.internal.backends.api_model import (
+    AttributeType,
+    AttributeWithProperties,
+    LeaderboardEntry,
+)
+from neptune.internal.backends.hosted_client import DEFAULT_REQUEST_KWARGS
+
+if TYPE_CHECKING:
+    from neptune.internal.backends.swagger_client_wrapper import SwaggerClientWrapper
+    from neptune.internal.id_formats import UniqueId
+
+
+SUPPORTED_ATTRIBUTE_TYPES = {item.value for item in AttributeType}
+
+
+def get_single_page(
+    *,
+    client: "SwaggerClientWrapper",
+    project_id: "UniqueId",
+    query_params: Dict[str, Any],
+    attributes_filter: Dict[str, Any],
+    limit: int,
+    offset: int,
+    types: Optional[Iterable[str]] = None,
+) -> List[Any]:
+    params = {
+        "projectIdentifier": project_id,
+        "type": types,
+        "params": {
+            **query_params,
+            **attributes_filter,
+            "pagination": {"limit": limit, "offset": offset},
+        },
+    }
+
+    request_options = DEFAULT_REQUEST_KWARGS.get("_request_options", {})
+    request_config = RequestConfig(request_options, True)
+    request_params = construct_request(client.api.searchLeaderboardEntries, request_options, **params)
+
+    http_client = client.swagger_spec.http_client
+
+    result = (
+        http_client.request(request_params, operation=None, request_config=request_config)
+        .response()
+        .incoming_response.json()
+    )
+
+    return list(result.get("entries", []))
+
+
+def to_leaderboard_entry(*, entry: Dict[str, Any]) -> LeaderboardEntry:
+    return LeaderboardEntry(
+        id=entry["experimentId"],
+        attributes=[
+            AttributeWithProperties(
+                path=attr["name"],
+                type=AttributeType(attr["type"]),
+                properties=attr.__getitem__(f"{attr['type']}Properties"),
+            )
+            for attr in entry["attributes"]
+            if attr["type"] in SUPPORTED_ATTRIBUTE_TYPES
+        ],
+    )
+
+
+def iter_over_pages(
+    *, iter_once: Callable[..., List[Any]], step: int, max_server_offset: int = 10000
+) -> Generator[Any, None, None]:
+    previous_items = None
+    num_of_collected_items = 0
+
+    while (previous_items is None or len(previous_items) >= step) and num_of_collected_items < max_server_offset:
+        previous_items = iter_once(
+            limit=min(step, max_server_offset - num_of_collected_items), offset=num_of_collected_items
+        )
+        num_of_collected_items += len(previous_items)
+        yield from previous_items

src/neptune/internal/backends/hosted_neptune_backend.py

@@ -20,6 +20,7 @@
 import os
 import re
 import typing
+from functools import partial
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -39,6 +40,11 @@
 )
 
 from neptune.api.dtos import FileEntry
+from neptune.api.searching_entries import (
+    get_single_page,
+    iter_over_pages,
+    to_leaderboard_entry,
+)
 from neptune.common.backends.utils import with_api_exceptions_handler
 from neptune.common.exceptions import (
     ClientHttpError,
@@ -67,7 +73,6 @@
     ArtifactAttribute,
     Attribute,
     AttributeType,
-    AttributeWithProperties,
     BoolAttribute,
     DatetimeAttribute,
     FileAttribute,
@@ -1022,43 +1027,27 @@ def search_leaderboard_entries(
         query: Optional[NQLQuery] = None,
         columns: Optional[Iterable[str]] = None,
     ) -> List[LeaderboardEntry]:
-        if query:
-            query_params = {"query": {"query": str(query)}}
-        else:
-            query_params = {}
-        if columns:
-            attributes_filter = {"attributeFilters": [{"path": column} for column in columns]}
-        else:
-            attributes_filter = {}
+        step_size = int(os.getenv(NEPTUNE_FETCH_TABLE_STEP_SIZE, "100"))
 
-        def get_portion(limit, offset):
-            return (
-                self.leaderboard_client.api.searchLeaderboardEntries(
-                    projectIdentifier=project_id,
-                    type=list(map(lambda container_type: container_type.to_api(), types)),
-                    params={
-                        **query_params,
-                        **attributes_filter,
-                        "pagination": {"limit": limit, "offset": offset},
-                    },
-                    **DEFAULT_REQUEST_KWARGS,
-                )
-                .response()
-                .result.entries
-            )
-
-        def to_leaderboard_entry(entry) -> LeaderboardEntry:
-            supported_attribute_types = {item.value for item in AttributeType}
-            attributes: List[AttributeWithProperties] = []
-            for attr in entry.attributes:
-                if attr.type in supported_attribute_types:
-                    properties = attr.__getitem__("{}Properties".format(attr.type))
-                    attributes.append(AttributeWithProperties(attr.name, AttributeType(attr.type), properties))
-            return LeaderboardEntry(entry.experimentId, attributes)
+        types_filter = list(map(lambda container_type: container_type.to_api(), types)) if types else None
+        query_params = {"query": {"query": str(query)}} if query else {}
+        attributes_filter = {"attributeFilters": [{"path": column} for column in columns]} if columns else {}
 
         try:
-            step_size = int(os.getenv(NEPTUNE_FETCH_TABLE_STEP_SIZE, "100"))
-            return [to_leaderboard_entry(e) for e in self._get_all_items(get_portion, step=step_size)]
+            return [
+                to_leaderboard_entry(entry=entry)
+                for entry in iter_over_pages(
+                    iter_once=partial(
+                        get_single_page,
+                        client=self.leaderboard_client,
+                        project_id=project_id,
+                        types=types_filter,
+                        query_params=query_params,
+                        attributes_filter=attributes_filter,
+                    ),
+                    step=step_size,
+                )
+            ]
         except HTTPNotFound:
             raise ProjectNotFound(project_id)
 
@@ -1084,13 +1073,3 @@ def get_model_version_url(
     ) -> str:
         base_url = self.get_display_address()
         return f"{base_url}/{workspace}/{project_name}/m/{model_id}/v/{sys_id}"
-
-    @staticmethod
-    def _get_all_items(get_portion, step):
-        max_server_offset = 10000
-        items = []
-        previous_items = None
-        while (previous_items is None or len(previous_items) >= step) and len(items) < max_server_offset:
-            previous_items = get_portion(limit=step, offset=len(items))
-            items += previous_items
-        return items

src/neptune/metadata_containers/metadata_containers_table.py

@@ -69,31 +69,31 @@ def get_attribute_value(self, path: str) -> Any:
             if attr.path == path:
                 _type = attr.type
                 if _type == AttributeType.RUN_STATE:
-                    return RunState.from_api(attr.properties.value).value
+                    return RunState.from_api(attr.properties.get("value")).value
                 if _type in (
                     AttributeType.FLOAT,
                     AttributeType.INT,
                     AttributeType.BOOL,
                     AttributeType.STRING,
                     AttributeType.DATETIME,
                 ):
-                    return attr.properties.value
+                    return attr.properties.get("value")
                 if _type == AttributeType.FLOAT_SERIES or _type == AttributeType.STRING_SERIES:
-                    return attr.properties.last
+                    return attr.properties.get("last")
                 if _type == AttributeType.IMAGE_SERIES:
                     raise MetadataInconsistency("Cannot get value for image series.")
                 if _type == AttributeType.FILE:
                     raise MetadataInconsistency("Cannot get value for file attribute. Use download() instead.")
                 if _type == AttributeType.FILE_SET:
                     raise MetadataInconsistency("Cannot get value for file set attribute. Use download() instead.")
                 if _type == AttributeType.STRING_SET:
-                    return set(attr.properties.values)
+                    return set(attr.properties.get("values"))
                 if _type == AttributeType.GIT_REF:
-                    return attr.properties.commit.commitId
+                    return attr.properties.get("commit", {}).get("commitId")
                 if _type == AttributeType.NOTEBOOK_REF:
-                    return attr.properties.notebookName
+                    return attr.properties.get("notebookName")
                 if _type == AttributeType.ARTIFACT:
-                    return attr.properties.hash
+                    return attr.properties.get("hash")
                 logger.error(
                     "Attribute type %s not supported in this version, yielding None. Recommended client upgrade.",
                     _type,
@@ -183,29 +183,29 @@ def make_attribute_value(
             _type = attribute.type
             _properties = attribute.properties
             if _type == AttributeType.RUN_STATE:
-                return RunState.from_api(_properties.value).value
+                return RunState.from_api(_properties.get("value")).value
             if _type in (
                 AttributeType.FLOAT,
                 AttributeType.INT,
                 AttributeType.BOOL,
                 AttributeType.STRING,
                 AttributeType.DATETIME,
             ):
-                return _properties.value
+                return _properties.get("value")
             if _type == AttributeType.FLOAT_SERIES or _type == AttributeType.STRING_SERIES:
-                return _properties.last
+                return _properties.get("last")
             if _type == AttributeType.IMAGE_SERIES:
                 return None
             if _type == AttributeType.FILE or _type == AttributeType.FILE_SET:
                 return None
             if _type == AttributeType.STRING_SET:
-                return ",".join(_properties.values)
+                return ",".join(_properties.get("values"))
             if _type == AttributeType.GIT_REF:
-                return _properties.commit.commitId
+                return _properties.get("commit", {}).get("commitId")
             if _type == AttributeType.NOTEBOOK_REF:
-                return _properties.notebookName
+                return _properties.get("notebookName")
             if _type == AttributeType.ARTIFACT:
-                return _properties.hash
+                return _properties.get("hash")
             logger.error(
                 "Attribute type %s not supported in this version, yielding None. Recommended client upgrade.",
                 _type,