Repository: SeldonIO/MLServer

Issue:
Compress requests / responses
Leverage `gzip` to compress requests / responses.


Discussion:
I'd be cautious about how this is implemented - for small payloads, the compression/decompression time might be a more significant cost than the network's speed.
Also for small payloads or ones without much repetition in their content, the compressed size might be larger than its uncompressed equivalent because of the dictionary's overheads.

There are probably some heuristics about what dictionary size to use and when to apply compression.

Resolving patch:
diff --git a/mlserver/rest/app.py b/mlserver/rest/app.py
--- a/mlserver/rest/app.py
+++ b/mlserver/rest/app.py
@@ -3,6 +3,7 @@
 from fastapi.responses import Response as FastAPIResponse
 from fastapi.routing import APIRoute as FastAPIRoute
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi.middleware.gzip import GZipMiddleware
 from starlette_exporter import PrometheusMiddleware
 
 from .endpoints import Endpoints, ModelRepositoryEndpoints
@@ -103,6 +104,7 @@ def create_app(
         exception_handlers=_EXCEPTION_HANDLERS,  # type: ignore
     )
 
+    app.add_middleware(GZipMiddleware)
     if settings.cors_settings is not None:
         app.add_middleware(
             CORSMiddleware,


Test patch:
diff --git a/tests/fixtures.py b/tests/fixtures.py
--- a/tests/fixtures.py
+++ b/tests/fixtures.py
@@ -1,4 +1,8 @@
 import asyncio
+import random
+import string
+
+from typing import Dict
 
 from mlserver import MLModel
 from mlserver.types import InferenceRequest, InferenceResponse, Parameters
@@ -12,6 +16,12 @@ class SumModel(MLModel):
     async def my_payload(self, payload: list) -> int:
         return sum(payload)
 
+    @custom_handler(rest_path="/custom-endpoint-with-long-response")
+    async def long_response_endpoint(self, length: int) -> Dict[str, str]:
+        alphabet = string.ascii_lowercase
+        response = "".join(random.choice(alphabet) for i in range(length))
+        return {"foo": response}
+
     async def predict(self, payload: InferenceRequest) -> InferenceResponse:
         decoded = self.decode(payload.inputs[0])
         total = decoded.sum(axis=1, keepdims=True)
diff --git a/tests/handlers/test_custom.py b/tests/handlers/test_custom.py
--- a/tests/handlers/test_custom.py
+++ b/tests/handlers/test_custom.py
@@ -15,5 +15,7 @@ def test_custom_handler(sum_model: SumModel, custom_handler: CustomHandler):
 def test_get_custom_handlers(sum_model: SumModel, custom_handler: CustomHandler):
     handlers = get_custom_handlers(sum_model)
 
-    assert len(handlers) == 1
-    assert handlers[0] == (custom_handler, sum_model.my_payload)
+    # NOTE: The first handler is the "test" handler used to test large response
+    # payloads (e.g. to validate compression is working as expected)
+    assert len(handlers) == 2
+    assert handlers[1] == (custom_handler, sum_model.my_payload)
diff --git a/tests/rest/test_custom.py b/tests/rest/test_custom.py
--- a/tests/rest/test_custom.py
+++ b/tests/rest/test_custom.py
@@ -6,3 +6,15 @@ async def test_custom_handler(rest_client, inference_request: InferenceRequest):
 
     assert response.status_code == 200
     assert response.json() == 10
+
+
+async def test_gzip_compression(rest_client, inference_request, sum_model_settings):
+    endpoint = "/custom-endpoint-with-long-response"
+    response = await rest_client.post(
+        endpoint, params={"length": 1000}, headers={"accept-encoding": "gzip"}
+    )
+
+    assert response.status_code == 200
+    assert "content-encoding" in response.headers
+    assert response.headers["content-encoding"] == "gzip"
+    assert "foo" in response.json()
diff --git a/tests/rest/test_endpoints.py b/tests/rest/test_endpoints.py
--- a/tests/rest/test_endpoints.py
+++ b/tests/rest/test_endpoints.py
@@ -1,6 +1,11 @@
 import pytest
 
-from mlserver import types, __version__
+from mlserver import __version__
+from mlserver.types import (
+    InferenceResponse,
+    MetadataServerResponse,
+    MetadataModelResponse,
+)
 
 
 async def test_live(rest_client):
@@ -28,7 +33,7 @@ async def test_metadata(rest_client):
     endpoint = "/v2"
     response = await rest_client.get(endpoint)
 
-    metadata = types.MetadataServerResponse.parse_obj(response.json())
+    metadata = MetadataServerResponse.parse_obj(response.json())
 
     assert metadata.name == "mlserver"
     assert metadata.version == __version__
@@ -39,7 +44,7 @@ async def test_model_metadata(rest_client, sum_model_settings):
     endpoint = f"v2/models/{sum_model_settings.name}"
     response = await rest_client.get(endpoint)
 
-    metadata = types.MetadataModelResponse.parse_obj(response.json())
+    metadata = MetadataModelResponse.parse_obj(response.json())
 
     assert metadata.name == sum_model_settings.name
     assert metadata.platform == sum_model_settings.platform
@@ -58,7 +63,7 @@ async def test_infer(rest_client, inference_request, model_name, model_version):
 
     assert response.status_code == 200
 
-    prediction = types.InferenceResponse.parse_obj(response.json())
+    prediction = InferenceResponse.parse_obj(response.json())
     assert len(prediction.outputs) == 1
     assert prediction.outputs[0].data.__root__ == [6]