Repository: OpenMined/PySyft

Issue:
relu and relu_deriv require serializable workers
**This is a blocker for having [EMLaaS implemented](https://github.com/OpenMined/Grid/issues/203)**.

Relu and relu_derive require serializable workers. Apparently these operations require serializable workers in order to execute `.share` with pointer tensors. GridWorkers are currently not serializable.

Discussion:

Patch:
diff --git a/syft/frameworks/torch/crypto/securenn.py b/syft/frameworks/torch/crypto/securenn.py
index 5c0ba54ae5..8487f6c7bd 100644
--- a/syft/frameworks/torch/crypto/securenn.py
+++ b/syft/frameworks/torch/crypto/securenn.py
@@ -438,7 +438,7 @@ def relu(a_sh):
 # 2**62 would almost always lead to overflow).
 def division(x_sh, y_sh, bit_len_max=Q_BITS // 2):
     """ Performs division of encrypted numbers
-    
+
     Args:
         x_sh, y_sh (AdditiveSharingTensor): the private tensors on which the op applies
         bit_len_max: the number of bits needed to represent the highest value in the tensors
@@ -493,10 +493,10 @@ def division(x_sh, y_sh, bit_len_max=Q_BITS // 2):
 def maxpool(x_sh):
     """ Compute MaxPool: returns fresh shares of the max value in the input tensor
     and the index of this value in the flattened tensor
-    
+
     Args:
         x_sh (AdditiveSharingTensor): the private tensor on which the op applies
-        
+
     Returns:
         maximum value as an AdditiveSharingTensor
         index of this value in the flattened tensor as an AdditiveSharingTensor
diff --git a/syft/messaging/__init__.py b/syft/messaging/__init__.py
index 2c6db74f08..0728d01b43 100644
--- a/syft/messaging/__init__.py
+++ b/syft/messaging/__init__.py
@@ -1,10 +1,10 @@
+from syft.messaging.message import ForceObjectDeleteMessage
+from syft.messaging.message import GetShapeMessage
+from syft.messaging.message import IsNoneMessage
 from syft.messaging.message import Message
 from syft.messaging.message import Operation
 from syft.messaging.message import ObjectMessage
 from syft.messaging.message import ObjectRequestMessage
-from syft.messaging.message import IsNoneMessage
-from syft.messaging.message import GetShapeMessage
-from syft.messaging.message import ForceObjectDeleteMessage
 from syft.messaging.message import SearchMessage
 
 from syft.messaging.plan import Plan
diff --git a/syft/serde/serde.py b/syft/serde/serde.py
index ca791fd500..5d76bc5656 100644
--- a/syft/serde/serde.py
+++ b/syft/serde/serde.py
@@ -33,6 +33,7 @@ If different compressions are required, the worker can override the function app
 """
 from collections import OrderedDict
 
+import inspect
 import lz4
 from lz4 import (  # noqa: F401
     frame,
@@ -52,7 +53,7 @@ from syft.frameworks.torch.tensors.interpreters import AutogradTensor
 from syft.generic import pointers
 from syft.serde.native_serde import MAP_NATIVE_SIMPLIFIERS_AND_DETAILERS
 from syft.workers import AbstractWorker
-from syft.workers import VirtualWorker
+from syft.workers import BaseWorker
 
 from syft.exceptions import CompressionNotFoundException
 from syft.exceptions import GetNotPermittedError
@@ -86,7 +87,7 @@ OBJ_SIMPLIFIER_AND_DETAILERS = [
     pointers.PointerTensor,
     pointers.ObjectWrapper,
     TrainConfig,
-    VirtualWorker,
+    BaseWorker,
     AutogradTensor,
     messaging.Message,
     messaging.Operation,
@@ -99,7 +100,7 @@ OBJ_SIMPLIFIER_AND_DETAILERS = [
 ]
 
 # If an object implements its own force_simplify and force_detail functions it should be stored in this list
-OBJ_FORCE_FULL_SIMPLIFIER_AND_DETAILERS = [VirtualWorker]
+OBJ_FORCE_FULL_SIMPLIFIER_AND_DETAILERS = [BaseWorker]
 
 # For registering syft objects with custom simplify and detail methods
 EXCEPTION_SIMPLIFIER_AND_DETAILERS = [GetNotPermittedError, ResponseSignatureError]
@@ -116,27 +117,52 @@ scheme_to_bytes = {
 
 ## SECTION: High Level Simplification Router
 def _force_full_simplify(obj: object) -> object:
-    """To force a full simplify genrally if the usual _simplify is not suitable.
+    """To force a full simplify generally if the usual _simplify is not suitable.
+
+    If we can not full simplify a object we simplify it as usual instead.
 
     Args:
-        The object
+        obj: The object.
 
     Returns:
-        The simplified object
+        The simplified object.
     """
+    # check to see if there is a full simplifier
+    # for this type. If there is, return the full simplified object.
     current_type = type(obj)
-
     if current_type in forced_full_simplifiers:
-        left = forced_full_simplifiers[current_type][0]
-
-        right = forced_full_simplifiers[current_type][1]
-        right = right(obj)
-
-        result = (left, right)
+        result = (
+            forced_full_simplifiers[current_type][0],
+            forced_full_simplifiers[current_type][1](obj),
+        )
+    # If we already tried to find a full simplifier for this type but failed, we should
+    # simplify it instead.
+    elif current_type in no_full_simplifiers_found:
+        return _simplify(obj)
     else:
-        result = _simplify(obj)
-
-    return result
+        # If the object type is not in forced_full_simplifiers,
+        # we check the classes that this object inherits from.
+        # `inspect.getmro` give us all types this object inherits
+        # from, including `type(obj)`. We can skip the type of the
+        # object because we already tried this in the
+        # previous step.
+        classes_inheritance = inspect.getmro(type(obj))[1:]
+
+        for inheritance_type in classes_inheritance:
+            if inheritance_type in forced_full_simplifiers:
+                # Store the inheritance_type in forced_full_simplifiers so next
+                # time we see this type serde will be faster.
+                forced_full_simplifiers[current_type] = forced_full_simplifiers[inheritance_type]
+                result = (
+                    forced_full_simplifiers[current_type][0],
+                    forced_full_simplifiers[current_type][1](obj),
+                )
+                return result
+
+        # If there is not a full_simplifier for this
+        # object, then we simplify it.
+        no_full_simplifiers_found.add(current_type)
+        return _simplify(obj)
 
 
 ## SECTION: dinamically generate simplifiers and detailers
@@ -184,6 +210,9 @@ def _generate_simplifiers_and_detailers():
 
 
 simplifiers, forced_full_simplifiers, detailers = _generate_simplifiers_and_detailers()
+# Store types that are not simplifiable (int, float, None) so we
+# can ignore them during serialization.
+no_simplifiers_found, no_full_simplifiers_found = set(), set()
 
 
 ## SECTION:  High Level Public Functions (these are the ones you use)
@@ -422,28 +451,50 @@ def _simplify(obj: object) -> object:
     being sent.
 
     Args:
-        obj: an object which may need to be simplified.
+        obj: An object which may need to be simplified.
 
     Returns:
-        obj: an simple Python object which msgpack can serialize.
+        An simple Python object which msgpack can serialize.
 
     Raises:
         ValueError: if `move_this` or `in_front_of_that` are not both single ASCII
         characters.
     """
-    try:
-        # check to see if there is a simplifier
-        # for this type. If there is, run return
-        # the simplified object
-        current_type = type(obj)
+
+    # Check to see if there is a simplifier
+    # for this type. If there is, return the simplified object.
+    current_type = type(obj)
+    if current_type in simplifiers:
         result = (simplifiers[current_type][0], simplifiers[current_type][1](obj))
         return result
 
-    except KeyError:
+    # If we already tried to find a simplifier for this type but failed, we should
+    # just return the object as it is.
+    elif current_type in no_simplifiers_found:
+        return obj
+
+    else:
+        # If the object type is not in simplifiers,
+        # we check the classes that this object inherits from.
+        # `inspect.getmro` give us all types this object inherits
+        # from, including `type(obj)`. We can skip the type of the
+        # object because we already tried this in the
+        # previous step.
+        classes_inheritance = inspect.getmro(type(obj))[1:]
+
+        for inheritance_type in classes_inheritance:
+            if inheritance_type in simplifiers:
+                # Store the inheritance_type in simplifiers so next time we see this type
+                # serde will be faster.
+                simplifiers[current_type] = simplifiers[inheritance_type]
+                result = (simplifiers[current_type][0], simplifiers[current_type][1](obj))
+                return result
+
         # if there is not a simplifier for this
         # object, then the object is already a
         # simple python object and we can just
-        # return it
+        # return it.
+        no_simplifiers_found.add(current_type)
         return obj
 
 
diff --git a/syft/workers/base.py b/syft/workers/base.py
index 55b269372c..7d877983eb 100644
--- a/syft/workers/base.py
+++ b/syft/workers/base.py
@@ -19,6 +19,7 @@ from syft.exceptions import GetNotPermittedError
 from syft.exceptions import WorkerNotFoundException
 from syft.exceptions import ResponseSignatureError
 
+
 # this if statement avoids circular imports between base.py and pointer.py
 if TYPE_CHECKING:
     from syft.generic import pointers
@@ -550,7 +551,7 @@ class BaseWorker(AbstractWorker, ObjectStorage):
 
     def get_worker(
         self, id_or_worker: Union[str, int, "BaseWorker"], fail_hard: bool = False
-    ) -> Union[str, int]:
+    ) -> Union[str, int, AbstractWorker]:
         """Returns the worker id or instance.
 
         Allows for resolution of worker ids to workers to happen automatically
@@ -591,23 +592,31 @@ class BaseWorker(AbstractWorker, ObjectStorage):
             >>> me.get_worker(bob)
             <syft.core.workers.virtual.VirtualWorker id:bob>
         """
-
         if isinstance(id_or_worker, bytes):
             id_or_worker = str(id_or_worker, "utf-8")
 
-        if isinstance(id_or_worker, (str, int)):
-            if id_or_worker in self._known_workers:
-                return self._known_workers[id_or_worker]
-            else:
-                if fail_hard:
-                    raise WorkerNotFoundException
-                logger.warning("Worker %s couldn't recognize worker %s", self.id, id_or_worker)
-                return id_or_worker
+        if isinstance(id_or_worker, str) or isinstance(id_or_worker, int):
+            return self._get_worker_based_on_id(id_or_worker, fail_hard=fail_hard)
         else:
-            if id_or_worker.id not in self._known_workers:
-                self.add_worker(id_or_worker)
+            return self._get_worker(id_or_worker)
+
+    def _get_worker(self, worker: AbstractWorker):
+        if worker.id not in self._known_workers:
+            self.add_worker(worker)
+        return worker
+
+    def _get_worker_based_on_id(self, worker_id: Union[str, int], fail_hard: bool = False):
+        # A worker should always know itself
+        if worker_id == self.id:
+            return self
+
+        worker = self._known_workers.get(worker_id, worker_id)
 
-        return id_or_worker
+        if worker == worker_id:
+            if fail_hard:
+                raise WorkerNotFoundException
+            logger.warning("Worker %s couldn't recognize worker %s", self.id, worker_id)
+        return worker
 
     def add_worker(self, worker: "BaseWorker"):
         """Adds a single worker.
@@ -867,3 +876,44 @@ class BaseWorker(AbstractWorker, ObjectStorage):
         return messaging.Message(
             codes.MSGTYPE.CMD, [[command_name, command_owner, args, kwargs], return_ids]
         )
+
+    @staticmethod
+    def simplify(worker: AbstractWorker) -> tuple:
+        return (sy.serde._simplify(worker.id),)
+
+    @staticmethod
+    def detail(worker: AbstractWorker, worker_tuple: tuple) -> Union["AbstractWorker", int, str]:
+        """
+        This function reconstructs a PlanPointer given it's attributes in form of a tuple.
+
+        Args:
+            worker: the worker doing the deserialization
+            plan_pointer_tuple: a tuple holding the attributes of the PlanPointer
+        Returns:
+            A worker id or worker instance.
+        """
+        worker_id = sy.serde._detail(worker, worker_tuple[0])
+
+        referenced_worker = worker.get_worker(worker_id)
+
+        return referenced_worker
+
+    @staticmethod
+    def force_simplify(worker: AbstractWorker) -> tuple:
+        return (sy.serde._simplify(worker.id), sy.serde._simplify(worker._objects), worker.auto_add)
+
+    @staticmethod
+    def force_detail(worker: AbstractWorker, worker_tuple: tuple) -> tuple:
+        worker_id, _objects, auto_add = worker_tuple
+        worker_id = sy.serde._detail(worker, worker_id)
+
+        result = sy.VirtualWorker(sy.hook, worker_id, auto_add=auto_add)
+        _objects = sy.serde._detail(worker, _objects)
+        result._objects = _objects
+
+        # make sure they weren't accidentally double registered
+        for _, obj in _objects.items():
+            if obj.id in worker._objects:
+                del worker._objects[obj.id]
+
+        return result
diff --git a/syft/workers/virtual.py b/syft/workers/virtual.py
index 3da2fd1bc3..5627bdada4 100644
--- a/syft/workers/virtual.py
+++ b/syft/workers/virtual.py
@@ -11,46 +11,3 @@ class VirtualWorker(BaseWorker, FederatedClient):
 
     def _recv_msg(self, message: bin) -> bin:
         return self.recv_msg(message)
-
-    @staticmethod
-    def force_simplify(worker: AbstractWorker) -> tuple:
-        return (sy.serde._simplify(worker.id), sy.serde._simplify(worker._objects), worker.auto_add)
-
-    @staticmethod
-    def force_detail(worker: AbstractWorker, worker_tuple: tuple) -> tuple:
-        worker_id, _objects, auto_add = worker_tuple
-        worker_id = sy.serde._detail(worker, worker_id)
-
-        result = sy.VirtualWorker(sy.hook, worker_id, auto_add=auto_add)
-        _objects = sy.serde._detail(worker, _objects)
-        result._objects = _objects
-
-        # make sure they weren't accidentally double registered
-        for _, obj in _objects.items():
-            if obj.id in worker._objects:
-                del worker._objects[obj.id]
-
-        return result
-
-    @staticmethod
-    def simplify(worker: AbstractWorker) -> tuple:
-        return (sy.serde._simplify(worker.id),)
-
-    @staticmethod
-    def detail(worker: AbstractWorker, worker_tuple: tuple) -> "pointers.PointerTensor":
-        """
-        This function reconstructs a PlanPointer given it's attributes in form of a tuple.
-
-        Args:
-            worker: the worker doing the deserialization
-            plan_pointer_tuple: a tuple holding the attributes of the PlanPointer
-        Returns:
-            PointerTensor: a PointerTensor
-        Examples:
-            ptr = detail(data)
-        """
-        worker_id = sy.serde._detail(worker, worker_tuple[0])
-
-        referenced_worker = worker.get_worker(worker_id)
-
-        return referenced_worker


Test patch:
diff --git a/test/test_serde.py b/test/test_serde.py
index 833881db17..cb04c1e248 100644
--- a/test/test_serde.py
+++ b/test/test_serde.py
@@ -707,3 +707,13 @@ def test_serde_object_wrapper_traced_module():
 
     assert (pred_before == pred_after).all()
     assert obj_wrapper.id == obj_wrapper_received.id
+
+
+def test_no_simplifier_found():
+    """Test that types that can not be simplified are cached."""
+    # Clean cache.
+    serde.no_simplifiers_found = set()
+    x = 1.3
+    assert type(x) not in serde.no_simplifiers_found
+    _ = serde._simplify(x)
+    assert type(x) in serde.no_simplifiers_found