Repository: modin-project/modin

Issue:
Improve RPyC connection performance
Currently (as of #1818) RPyC is being used as the main communication channel between local Modin and scaled-out to the cluster.
It does a lot of excessive calls, and each call costs at least as much time as ping to spawned cluster (because most of the calls need to be synchronous).

We need to make less calls and cache whatever is safe.


Discussion:

Resolving patch:
diff --git a/modin/data_management/factories.py b/modin/data_management/factories.py
--- a/modin/data_management/factories.py
+++ b/modin/data_management/factories.py
@@ -241,8 +241,7 @@ def __getattr__(self, name):
                     except KeyError:
 
                         def wrap(*a, _original=getattr(self.__io_cls, name), **kw):
-                            a = tuple(self.__conn.deliver(x) for x in a)
-                            kw = {k: self.__conn.deliver(v) for k, v in kw.items()}
+                            a, kw = self.__conn.deliver(a, kw)
                             return _original(*a, **kw)
 
                         self.__wrappers[name] = wrap
diff --git a/modin/experimental/cloud/connection.py b/modin/experimental/cloud/connection.py
--- a/modin/experimental/cloud/connection.py
+++ b/modin/experimental/cloud/connection.py
@@ -106,11 +106,12 @@ def __try_connect(self):
         from .rpyc_proxy import WrappingService
 
         try:
-            self.__connection = rpyc.connect(
-                "127.0.0.1",
-                self.rpyc_port,
+            stream = rpyc.SocketStream.connect(
+                host="127.0.0.1", port=self.rpyc_port, nodelay=True, keepalive=True
+            )
+            self.__connection = rpyc.connect_stream(
+                stream,
                 WrappingService,
-                keepalive=True,
                 config={"sync_request_timeout": RPYC_REQUEST_TIMEOUT},
             )
         except (ConnectionRefusedError, EOFError):
diff --git a/modin/experimental/cloud/rpyc_proxy.py b/modin/experimental/cloud/rpyc_proxy.py
--- a/modin/experimental/cloud/rpyc_proxy.py
+++ b/modin/experimental/cloud/rpyc_proxy.py
@@ -12,34 +12,177 @@
 # governing permissions and limitations under the License.
 
 import types
+import collections
 
 import rpyc
 from rpyc.lib.compat import pickle
-from rpyc.core import netref
+from rpyc.lib import get_methods
+
+from rpyc.core import netref, AsyncResult, consts
 
 from . import get_connection
 from .meta_magic import _LOCAL_ATTRS, RemoteMeta, _KNOWN_DUALS
 
 
+def _batch_loads(items):
+    return tuple(pickle.loads(item) for item in items)
+
+
+def _tuplize(arg):
+    """turns any sequence or iterator into a flat tuple"""
+    return tuple(arg)
+
+
 class WrappingConnection(rpyc.Connection):
     def __init__(self, *a, **kw):
         super().__init__(*a, **kw)
-        self._remote_pickle_loads = None
+        self._remote_batch_loads = None
+        self._remote_cls_cache = {}
+        self._static_cache = collections.defaultdict(dict)
+        self._remote_dumps = None
+        self._remote_tuplize = None
+
+    def __wrap(self, local_obj):
+        while True:
+            # unwrap magic wrappers first; keep unwrapping in case it's a wrapper-in-a-wrapper
+            # this shouldn't usually happen, so this is mostly a safety net
+            try:
+                local_obj = object.__getattribute__(local_obj, "__remote_end__")
+            except AttributeError:
+                break
+        # do not pickle netrefs of our current connection, but do pickle those of other;
+        # example of this: an object made in _other_ remote context being pased to ours
+        if isinstance(local_obj, netref.BaseNetref) and local_obj.____conn__ is self:
+            return None
+        return bytes(pickle.dumps(local_obj))
+
+    def deliver(self, args, kw):
 
-    def deliver(self, local_obj):
         """
-        More caching version of rpyc.classic.deliver()
+        More efficient, batched version of rpyc.classic.deliver()
         """
-        try:
-            local_obj = object.__getattribute__(local_obj, "__remote_end__")
-        except AttributeError:
-            pass
-        if isinstance(local_obj, netref.BaseNetref) and local_obj.____conn__ is self:
-            return local_obj
-        return self._remote_pickle_loads(bytes(pickle.dumps(local_obj)))
+        pickled_args = [self.__wrap(arg) for arg in args]
+        pickled_kw = [(k, self.__wrap(v)) for (k, v) in kw.items()]
+
+        pickled = [i for i in pickled_args if i is not None] + [
+            v for (k, v) in pickled_kw if v is not None
+        ]
+        remote = iter(self._remote_batch_loads(tuple(pickled)))
+
+        delivered_args = []
+        for local_arg, pickled_arg in zip(args, pickled_args):
+            delivered_args.append(
+                next(remote) if pickled_arg is not None else local_arg
+            )
+        delivered_kw = {}
+        for k, pickled_v in pickled_kw:
+            delivered_kw[k] = next(remote) if pickled_v is not None else kw[k]
+
+        return tuple(delivered_args), delivered_kw
+
+    def obtain(self, remote):
+        while True:
+            try:
+                remote = object.__getattribute__(remote, "__remote_end__")
+            except AttributeError:
+                break
+        return pickle.loads(self._remote_dumps(remote))
+
+    def obtain_tuple(self, remote):
+        while True:
+            try:
+                remote = object.__getattribute__(remote, "__remote_end__")
+            except AttributeError:
+                break
+        return self._remote_tuplize(remote)
+
+    def sync_request(self, handler, *args):
+        """
+        Intercept outgoing synchronous requests from RPyC to add caching or
+        fulfilling them locally if possible to improve performance.
+        We should try to make as few remote calls as possible, because each
+        call adds up to latency.
+        """
+        if handler == consts.HANDLE_INSPECT:
+            # always inspect classes from modin, pandas and numpy locally,
+            # do not go to network for those
+            id_name = str(args[0][0])
+            if id_name.split(".", 1)[0] in ("modin", "pandas", "numpy"):
+                try:
+                    modobj = __import__(id_name)
+                    for subname in id_name.split(".")[1:]:
+                        modobj = getattr(modobj, subname)
+                except (ImportError, AttributeError):
+                    pass
+                else:
+                    return get_methods(netref.LOCAL_ATTRS, modobj)
+                modname, clsname = id_name.rsplit(".", 1)
+                try:
+                    modobj = __import__(modname)
+                    for subname in modname.split(".")[1:]:
+                        modobj = getattr(modobj, subname)
+                    clsobj = getattr(modobj, clsname)
+                except (ImportError, AttributeError):
+                    pass
+                else:
+                    return get_methods(netref.LOCAL_ATTRS, clsobj)
+        elif handler in (consts.HANDLE_GETATTR, consts.HANDLE_STR, consts.HANDLE_HASH):
+            if handler == consts.HANDLE_GETATTR:
+                obj, attr = args
+                key = (attr, handler)
+            else:
+                obj = args[0]
+                key = handler
+
+            if str(obj.____id_pack__[0]) in {"numpy", "numpy.dtype"}:
+                # always assume numpy attributes and numpy.dtype attributes are always the same;
+                # note that we're using RPyC id_pack as cache key, and it includes the name,
+                # class id and instance id, so this cache is unique to each instance of, say,
+                # numpy.dtype(), hence numpy.int16 and numpy.float64 got different caches.
+                cache = self._static_cache[obj.____id_pack__]
+                try:
+                    result = cache[key]
+                except KeyError:
+                    result = cache[key] = super().sync_request(handler, *args)
+                    if handler == consts.HANDLE_GETATTR:
+                        # save an entry in our cache telling that we get this attribute cached
+                        self._static_cache[result.____id_pack__]["__getattr__"] = True
+                return result
+
+        return super().sync_request(handler, *args)
+
+    def async_request(self, handler, *args, **kw):
+        """
+        Override async request handling to intercept outgoing deletion requests because we cache
+        certain things, and if we allow deletion of cached things our cache becomes stale.
+        We can clean the cache upon deletion, but it would increase cache misses a lot.
+
+        Also note that memory is not leaked forever, RPyC frees all of it upon disconnect.
+        """
+        if handler == consts.HANDLE_DEL:
+            obj, _ = args
+            if obj.____id_pack__ in self._static_cache:
+                # object is cached by us, so ignore the request or remote end dies and cache is suddenly stale;
+                # we shouldn't remove item from cache as it would reduce performance
+                res = AsyncResult(self)
+                res._is_ready = True  # simulate finished async request
+                return res
+        return super().async_request(handler, *args, **kw)
 
     def _netref_factory(self, id_pack):
-        result = super()._netref_factory(id_pack)
+        id_name, cls_id, inst_id = id_pack
+        id_name = str(id_name)
+        first = id_name.split(".", 1)[0]
+        if first in ("modin", "numpy", "pandas") and inst_id:
+            try:
+                cached_cls = self._remote_cls_cache[(id_name, cls_id)]
+            except KeyError:
+                result = super()._netref_factory(id_pack)
+                self._remote_cls_cache[(id_name, cls_id)] = type(result)
+            else:
+                result = cached_cls(self, id_pack)
+        else:
+            result = super()._netref_factory(id_pack)
         # try getting __real_cls__ from result.__class__ BUT make sure to
         # NOT get it from some parent class for result.__class__, otherwise
         # multiple wrappings happen
@@ -78,7 +221,13 @@ def _box(self, obj):
         return super()._box(obj)
 
     def _init_deliver(self):
-        self._remote_pickle_loads = self.modules["rpyc.lib.compat"].pickle.loads
+        self._remote_batch_loads = self.modules[
+            "modin.experimental.cloud.rpyc_proxy"
+        ]._batch_loads
+        self._remote_dumps = self.modules["rpyc.lib.compat"].pickle.dumps
+        self._remote_tuplize = self.modules[
+            "modin.experimental.cloud.rpyc_proxy"
+        ]._tuplize
 
 
 class WrappingService(rpyc.ClassicService):
@@ -167,6 +316,7 @@ def __prepare__(*args, **kw):
             type(obj).__dict__ (EXCLUDING parent classes) and then goes to proxy type.
             """
             namespace = type.__prepare__(*args, **kw)
+            namespace["__remote_methods__"] = {}
 
             # try computing overridden differently to allow subclassing one override from another
             no_override = set(_NO_OVERRIDE)
@@ -204,9 +354,15 @@ def __prepare__(*args, **kw):
                     ):
 
                         def method(_self, *_args, __method_name__=name, **_kw):
-                            return getattr(_self.__remote_end__, __method_name__)(
-                                *_args, **_kw
-                            )
+                            try:
+                                remote = _self.__remote_methods__[__method_name__]
+                            except KeyError:
+                                # use remote_cls.__getattr__ to force RPyC return us
+                                # a proxy for remote method call instead of its local wrapper
+                                _self.__remote_methods__[
+                                    __method_name__
+                                ] = remote = remote_cls.__getattr__(__method_name__)
+                            return remote(_self.__remote_end__, *_args, **_kw)
 
                         method.__name__ = name
                         namespace[name] = method
@@ -353,9 +509,16 @@ class DeliveringMixin:
     for method in methods:
 
         def wrapper(self, *args, __remote_conn__=conn, __method_name__=method, **kw):
-            args = tuple(__remote_conn__.deliver(x) for x in args)
-            kw = {k: __remote_conn__.deliver(v) for k, v in kw.items()}
-            return getattr(self.__remote_end__, __method_name__)(*args, **kw)
+            args, kw = __remote_conn__.deliver(args, kw)
+            cache = object.__getattribute__(self, "__remote_methods__")
+            try:
+                remote = cache[__method_name__]
+            except KeyError:
+                # see comments in ProxyMeta.__prepare__ on using remote_cls.__getattr__
+                cache[__method_name__] = remote = remote_cls.__getattr__(
+                    __method_name__
+                )
+            return remote(self.__remote_end__, *args, **kw)
 
         wrapper.__name__ = method
         setattr(mixin, method, wrapper)
@@ -396,10 +559,30 @@ def make_dataframe_wrapper(DataFrame):
     It makes DF.loc, DF.groupby() and other methods listed below deliver their
     arguments to remote end by value.
     """
+
+    from modin.pandas.series import Series
+
+    conn = get_connection()
+
+    class ObtainingItems:
+        def items(self):
+            return conn.obtain_tuple(self.__remote_end__.items())
+
+        def iteritems(self):
+            return conn.obtain_tuple(self.__remote_end__.iteritems())
+
+    ObtainingItems = _deliveringWrapper(Series, mixin=ObtainingItems)
+
+    class DataFrameOverrides(_prepare_loc_mixin()):
+        @property
+        def dtypes(self):
+            remote_dtypes = self.__remote_end__.dtypes
+            return ObtainingItems(__remote_end__=remote_dtypes)
+
     DeliveringDataFrame = _deliveringWrapper(
         DataFrame,
         ["groupby", "agg", "aggregate", "__getitem__", "astype", "drop", "merge"],
-        _prepare_loc_mixin(),
+        DataFrameOverrides,
         "DataFrame",
     )
     return DeliveringDataFrame


Test patch: