Add/compression methods (#125)

AHReccese · web-flow · commit 9aef6195ed52 · 2024-08-23T23:07:29.000-07:00
* refactorings + update `parse` to return python obj

* add support of `gzip`, `zlib`, `lzma` and `bz2` compressing methods

* add compressor to PyMiloClient init params

* add docstring to compile method

* rename `compile` to `encrypt_compress`

* add compressor to PyMiloServer init params

* `autopep8.sh` applied

* update `run_server` script to set compression method

* update scenarios to set compression method in PyMiloClient

* raise concrete Exception type

* combine enc + comp from client side

* major refactorings + run ml streaming tests per compression method

* `autopep8.sh` applied
diff --git a/pymilo/streaming/communicator.py b/pymilo/streaming/communicator.py
@@ -1,5 +1,6 @@
 # -*- coding: utf-8 -*-
 """PyMilo RESTFull Communication Mediums."""
+import json
 import uvicorn
 import requests
 from pydantic import BaseModel
@@ -178,9 +179,9 @@ def parse(self, body):
         :type body: str
         :return: the extracted decrypted version
         """
-        return self._ps._compressor.extract(
-            self._ps._encryptor.decrypt(
-                body
+        return json.loads(
+            self._ps._compressor.extract(
+                self._ps._encryptor.decrypt(body)
             )
         )
 
diff --git a/pymilo/streaming/compressor.py b/pymilo/streaming/compressor.py
@@ -1,6 +1,13 @@
 # -*- coding: utf-8 -*-
 """Implementations of Compressor interface."""
-from .interfaces import Compressor
+import gzip
+import zlib
+import lzma
+import bz2
+import json
+import base64
+from enum import Enum
+from pymilo.streaming.interfaces import Compressor
 
 
 class DummyCompressor(Compressor):
@@ -9,9 +16,99 @@ class DummyCompressor(Compressor):
     @staticmethod
     def compress(payload):
         """Compress the given payload in a dummy way, simply just return it (no compression applied)."""
-        return payload
+        return payload if isinstance(payload, str) else json.dumps(payload)
 
     @staticmethod
     def extract(payload):
         """Extract the given payload in a dummy way, simply just return it (no Extraction applied)."""
         return payload
+
+
+class GZIPCompressor(Compressor):
+    """GZIP implementation of the Compressor interface."""
+
+    @staticmethod
+    def compress(payload):
+        """Compress the given payload using gzip."""
+        if isinstance(payload, str):
+            data = payload.encode('utf-8')
+        else:
+            data = json.dumps(payload).encode('utf-8')
+        compressed_data = gzip.compress(data)
+        return base64.b64encode(compressed_data).decode('utf-8')
+
+    @staticmethod
+    def extract(payload):
+        """Extract the given payload using gzip."""
+        data = base64.b64decode(payload)
+        return gzip.decompress(data).decode('utf-8')
+
+
+class ZLIBCompressor(Compressor):
+    """ZLIB implementation of the Compressor interface."""
+
+    @staticmethod
+    def compress(payload):
+        """Compress the given payload using zlib."""
+        if isinstance(payload, str):
+            data = payload.encode('utf-8')
+        else:
+            data = json.dumps(payload).encode('utf-8')
+        compressed_data = zlib.compress(data)
+        return base64.b64encode(compressed_data).decode('utf-8')
+
+    @staticmethod
+    def extract(payload):
+        """Extract the given payload using zlib."""
+        data = base64.b64decode(payload)
+        return zlib.decompress(data).decode('utf-8')
+
+
+class LZMACompressor(Compressor):
+    """LZMA implementation of the Compressor interface."""
+
+    @staticmethod
+    def compress(payload):
+        """Compress the given payload using lzma."""
+        if isinstance(payload, str):
+            data = payload.encode('utf-8')
+        else:
+            data = json.dumps(payload).encode('utf-8')
+        compressed_data = lzma.compress(data)
+        return base64.b64encode(compressed_data).decode('utf-8')
+
+    @staticmethod
+    def extract(payload):
+        """Extract the given payload using lzma."""
+        data = base64.b64decode(payload)
+        return lzma.decompress(data).decode('utf-8')
+
+
+class BZ2Compressor(Compressor):
+    """BZ2 implementation of the Compressor interface."""
+
+    @staticmethod
+    def compress(payload):
+        """Compress the given payload using bz2."""
+        if isinstance(payload, str):
+            data = payload.encode('utf-8')
+        else:
+            data = json.dumps(payload).encode('utf-8')
+        compressed_data = bz2.compress(data)
+        return base64.b64encode(compressed_data).decode('utf-8')
+
+    @staticmethod
+    def extract(payload):
+        """Extract the given payload using bz2."""
+        data = base64.b64decode(payload)
+        return bz2.decompress(data).decode('utf-8')
+
+
+class Compression(Enum):
+    """Compression method used in end to end communication."""
+
+    NULL = DummyCompressor
+    GZIP = GZIPCompressor
+    ZLIB = ZLIBCompressor
+    LZMA = LZMACompressor
+    BZ2 = BZ2Compressor
diff --git a/pymilo/streaming/pymilo_client.py b/pymilo/streaming/pymilo_client.py
@@ -2,7 +2,7 @@
 """PyMiloClient for RESTFull Protocol."""
 from enum import Enum
 from .encryptor import DummyEncryptor
-from .compressor import DummyCompressor
+from .compressor import Compression
 from ..pymilo_obj import Export, Import
 from .param import PYMILO_CLIENT_INVALID_MODE, PYMILO_CLIENT_MODEL_SYNCHED, \
     PYMILO_CLIENT_LOCAL_MODEL_UPLOADED, PYMILO_CLIENT_LOCAL_MODEL_UPLOAD_FAILED, \
@@ -25,6 +25,7 @@ def __init__(
             self,
             model=None,
             mode=Mode.LOCAL,
+            compressor=Compression.NULL,
             server_url="http://127.0.0.1:8000",
     ):
         """
@@ -34,6 +35,8 @@ def __init__(
         :type model: Any
         :param mode: the mode in which PymiloClient should work, either LOCAL mode or DELEGATE
         :type mode: str (LOCAL|DELEGATE)
+        :param compressor: the compression method to be used in client-server communications
+        :type compressor: pymilo.streaming.compressor.Compression
         :param server_url: the url to which PyMilo Server listens
         :type server_url: str
         :return: an instance of the Pymilo PymiloClient class
@@ -42,10 +45,24 @@ def __init__(
         self._model_id = "0x_model_id"
         self._model = model
         self._mode = mode
-        self._compressor = DummyCompressor()
+        self._compressor = compressor.value
         self._encryptor = DummyEncryptor()
         self._communicator = RESTClientCommunicator(server_url)
 
+    def encrypt_compress(self, body):
+        """
+        Compress and Encrypt body payload.
+
+        :param body: body payload of the request
+        :type body: dict
+        :return: the compressed and encrypted version of the body payload
+        """
+        return self._encryptor.encrypt(
+            self._compressor.compress(
+                body
+            )
+        )
+
     def toggle_mode(self, mode=Mode.LOCAL):
         """
         Toggle the PyMiloClient mode, either from LOCAL to DELEGATE or vice versa.
@@ -63,10 +80,14 @@ def download(self):
 
         :return: None
         """
-        serialized_model = self._communicator.download({
-            "client_id": self._client_id,
-            "model_id": self._model_id
-        })
+        serialized_model = self._communicator.download(
+            self.encrypt_compress(
+                {
+                    "client_id": self._client_id,
+                    "model_id": self._model_id,
+                }
+            )
+        )
         if serialized_model is None:
             print(PYMILO_CLIENT_FAILED_TO_DOWNLOAD_REMOTE_MODEL)
             return
@@ -79,11 +100,15 @@ def upload(self):
 
         :return: None
         """
-        succeed = self._communicator.upload({
-            "client_id": self._client_id,
-            "model_id": self._model_id,
-            "model": Export(self._model).to_json(),
-        })
+        succeed = self._communicator.upload(
+            self.encrypt_compress(
+                {
+                    "client_id": self._client_id,
+                    "model_id": self._model_id,
+                    "model": Export(self._model).to_json(),
+                }
+            )
+        )
         if succeed:
             print(PYMILO_CLIENT_LOCAL_MODEL_UPLOADED)
         else:
@@ -106,14 +131,12 @@ def __getattr__(self, attribute):
         elif self._mode == Mode.DELEGATE:
             gdst = GeneralDataStructureTransporter()
             response = self._communicator.attribute_type(
-                self._encryptor.encrypt(
-                    self._compressor.compress(
-                        {
-                            "client_id": self._client_id,
-                            "model_id": self._model_id,
-                            "attribute": attribute,
-                        }
-                    )
+                self.encrypt_compress(
+                    {
+                        "client_id": self._client_id,
+                        "model_id": self._model_id,
+                        "attribute": attribute,
+                    }
                 )
             )
             if response["attribute type"] == "field":
@@ -130,10 +153,8 @@ def relayer(*args, **kwargs):
                 payload["args"] = gdst.serialize(payload, "args", None)
                 payload["kwargs"] = gdst.serialize(payload, "kwargs", None)
                 result = self._communicator.attribute_call(
-                    self._encryptor.encrypt(
-                        self._compressor.compress(
-                            payload
-                        )
+                    self.encrypt_compress(
+                        payload
                     )
                 )
                 return gdst.deserialize(result, "payload", None)
diff --git a/pymilo/streaming/pymilo_server.py b/pymilo/streaming/pymilo_server.py
@@ -1,8 +1,8 @@
 # -*- coding: utf-8 -*-
 """PyMiloServer for RESTFull protocol."""
 from ..pymilo_obj import Export, Import
+from .compressor import Compression
 from .encryptor import DummyEncryptor
-from .compressor import DummyCompressor
 from .communicator import RESTServerCommunicator
 from .param import PYMILO_SERVER_NON_EXISTENT_ATTRIBUTE
 from ..transporters.general_data_structure_transporter import GeneralDataStructureTransporter
@@ -11,16 +11,18 @@
 class PymiloServer:
     """Facilitate streaming the ML models."""
 
-    def __init__(self, port=8000):
+    def __init__(self, port=8000, compressor=Compression.NULL):
         """
         Initialize the Pymilo PymiloServer instance.
 
         :param port: the port to which PyMiloServer listens
         :type port: int
+        :param compressor: the compression method to be used in client-server communications
+        :type compressor: pymilo.streaming.compressor.Compression
         :return: an instance of the PymiloServer class
         """
         self._model = None
-        self._compressor = DummyCompressor()
+        self._compressor = compressor.value
         self._encryptor = DummyEncryptor()
         self._communicator = RESTServerCommunicator(ps=self, port=port)
 
diff --git a/tests/test_ml_streaming/run_server.py b/tests/test_ml_streaming/run_server.py
@@ -1,4 +1,15 @@
+import argparse
+from pymilo.streaming.compressor import Compression
 from pymilo.streaming.pymilo_server import PymiloServer
 
-communicator = PymiloServer()._communicator
-communicator.run()
+
+def main():
+    parser = argparse.ArgumentParser(description='Run the Pymilo server with a specified compression method.')
+    parser.add_argument('--compression', type=str, choices=['NULL', 'GZIP', 'ZLIB', 'LZMA', 'BZ2'], default='NULL',
+                        help='Specify the compression method (NULL, GZIP, ZLIB, LZMA, or BZ2). Default is NULL.')
+    args = parser.parse_args()
+    communicator = PymiloServer(compressor=Compression[args.compression])._communicator
+    communicator.run()
+
+if __name__ == '__main__':
+    main()
diff --git a/tests/test_ml_streaming/scenarios/scenario1.py b/tests/test_ml_streaming/scenarios/scenario1.py
@@ -1,11 +1,12 @@
 import numpy as np
 from sklearn.metrics import mean_squared_error
 from sklearn.linear_model import LinearRegression
+from pymilo.streaming.compressor import Compression
 from pymilo.streaming.pymilo_client import PymiloClient, Mode
 from pymilo.utils.data_exporter import prepare_simple_regression_datasets
 
 
-def scenario1():
+def scenario1(compression_method):
     # 1. create model in local
     # 2. train model in local
     # 3. calculate mse before streaming
@@ -20,7 +21,7 @@ def scenario1():
 
     # 2.
     linear_regression.fit(x_train, y_train)
-    client = PymiloClient(model=linear_regression, mode=Mode.LOCAL)
+    client = PymiloClient(model=linear_regression, mode=Mode.LOCAL, compressor=Compression[compression_method])
 
     # 3.
     result = client.predict(x_test)
diff --git a/tests/test_ml_streaming/scenarios/scenario2.py b/tests/test_ml_streaming/scenarios/scenario2.py
@@ -1,11 +1,12 @@
 import numpy as np
 from sklearn.metrics import mean_squared_error
 from sklearn.linear_model import LinearRegression
+from pymilo.streaming.compressor import Compression
 from pymilo.streaming.pymilo_client import PymiloClient, Mode
 from pymilo.utils.data_exporter import prepare_simple_regression_datasets
 
 
-def scenario2():
+def scenario2(compression_method):
     # 1. create model in local
     # 2. upload model to server
     # 3. train model in server
@@ -17,7 +18,7 @@ def scenario2():
     # 1.
     x_train, y_train, x_test, y_test = prepare_simple_regression_datasets()
     linear_regression = LinearRegression()
-    client = PymiloClient(model=linear_regression, mode=Mode.LOCAL)
+    client = PymiloClient(model=linear_regression, mode=Mode.LOCAL, compressor=Compression[compression_method])
 
     # 2.
     client.upload()
diff --git a/tests/test_ml_streaming/test_streaming.py b/tests/test_ml_streaming/test_streaming.py
@@ -6,26 +6,30 @@
 from scenarios.scenario1 import scenario1
 from scenarios.scenario2 import scenario2
 
-@pytest.fixture(scope="session")
-def prepare_server():
+@pytest.fixture(scope="session", params=["NULL", "GZIP", "ZLIB", "LZMA", "BZ2"])
+def prepare_server(request):
+    compression_method = request.param
     path = os.path.join(
         os.getcwd(),
         "tests",
         "test_ml_streaming",
-        "run_server.py"
+        "run_server.py",
         )
     server_proc = subprocess.Popen(
         [
             executable,
             path,
+            "--compression", compression_method
         ],
         )
     time.sleep(2)
-    yield server_proc
+    yield (server_proc, compression_method)
     server_proc.terminate()
 
 def test1(prepare_server):
-    assert scenario1() == 0
+    _, compression_method = prepare_server
+    assert scenario1(compression_method) == 0
 
 def test2(prepare_server):
-    assert scenario2() == 0
+    _, compression_method = prepare_server
+    assert scenario2(compression_method) == 0
