diff --git a/CHANGELOG.md b/CHANGELOG.md index d11eb391e..e3b779867 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,7 @@ ### Changes - Upgraded performance of sending series data to Neptune ([#1483](https://github.com/neptune-ai/neptune-client/pull/1483)) - +- Compress (gzip) request to server, when server support it ([#1476](https://github.com/neptune-ai/neptune-client/pull/1476)) ## neptune 1.8.1 diff --git a/src/neptune/internal/backends/api_model.py b/src/neptune/internal/backends/api_model.py index 9dd87195d..88660ee12 100644 --- a/src/neptune/internal/backends/api_model.py +++ b/src/neptune/internal/backends/api_model.py @@ -132,6 +132,7 @@ class ClientConfig: _missing_features: FrozenSet[str] version_info: VersionInfo multipart_config: MultipartConfig + gzip_upload: bool def has_feature(self, feature_name: str) -> bool: return feature_name not in self._missing_features @@ -162,6 +163,7 @@ def from_api_response(config) -> "ClientConfig": multipart_upload_config = MultipartConfig( min_chunk_size, max_chunk_size, max_chunk_count, max_single_part_size ) + gzip_upload = getattr(config, "gzipUpload", False) artifacts_config_obj = getattr(config, "artifacts", None) has_artifacts = getattr(artifacts_config_obj, "enabled", False) @@ -179,6 +181,7 @@ def from_api_response(config) -> "ClientConfig": _missing_features=frozenset(missing_features), version_info=VersionInfo.build(min_recommended, min_compatible, max_compatible), multipart_config=multipart_upload_config, + gzip_upload=gzip_upload, ) diff --git a/src/neptune/internal/backends/hosted_client.py b/src/neptune/internal/backends/hosted_client.py index 0a74e77ef..239a4eaf0 100644 --- a/src/neptune/internal/backends/hosted_client.py +++ b/src/neptune/internal/backends/hosted_client.py @@ -23,6 +23,7 @@ import os import platform +import zlib from typing import ( Dict, Tuple, @@ -31,6 +32,11 @@ import requests from bravado.http_client import HttpClient from bravado.requests_client import RequestsClient +from requests import ( + PreparedRequest, + Response, +) +from requests.adapters import HTTPAdapter from neptune.common.backends.utils import with_api_exceptions_handler from neptune.common.oauth import NeptuneAuthenticator @@ -48,6 +54,7 @@ verify_host_resolution, ) from neptune.internal.credentials import Credentials +from neptune.internal.utils.logger import logger from neptune.version import version as neptune_client_version BACKEND_SWAGGER_PATH = "/api/backend/swagger.json" @@ -66,6 +73,22 @@ } +class GzipAdapter(HTTPAdapter): + def send(self, request: PreparedRequest, stream: bool = False, **kw) -> Response: + if request.body is not None and not stream and request.headers.get("Content-Type", None) == "application/json": + try: + request_body = request.body if isinstance(request.body, bytes) else bytes(request.body, "utf-8") + gzip_compress = zlib.compressobj(zlib.Z_DEFAULT_COMPRESSION, zlib.DEFLATED, zlib.MAX_WBITS | 16) + compressed = gzip_compress.compress(request_body) + gzip_compress.flush() + request.prepare_body(compressed, None) + request.headers["Content-Encoding"] = "gzip" + except zlib.error: + logger.warning("Error on compressing request") + pass + + return super(GzipAdapter, self).send(request, stream, **kw) + + def _close_connections_on_fork(session: requests.Session): try: os.register_at_fork(before=session.close, after_in_child=session.close, after_in_parent=session.close) @@ -174,7 +197,11 @@ def create_backend_client(client_config: ClientConfig, http_client: HttpClient) @cache -def create_leaderboard_client(client_config: ClientConfig, http_client: HttpClient) -> SwaggerClientWrapper: +def create_leaderboard_client(client_config: ClientConfig, http_client: RequestsClient) -> SwaggerClientWrapper: + if client_config.gzip_upload: + http_client.session.mount("http://", GzipAdapter()) + http_client.session.mount("https://", GzipAdapter()) + return SwaggerClientWrapper( create_swagger_client( build_operation_url(client_config.api_url, LEADERBOARD_SWAGGER_PATH),