FEAT-modin-project#2058: Improve how remote factories are defined

Signed-off-by: Vasilij Litvinov <[email protected]>

modin/data_management/__init__.py

@@ -10,3 +10,20 @@
 # the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
 # ANY KIND, either express or implied. See the License for the specific language
 # governing permissions and limitations under the License.
+
+from . import factories
+
+
+def _get_remote_engines():
+    for name in dir(factories):
+        obj = getattr(factories, name)
+        if isinstance(obj, type) and issubclass(
+            obj, factories.ExperimentalRemoteFactory
+        ):
+            try:
+                yield obj.get_info().engine
+            except factories.NotRealFactory:
+                pass
+
+
+REMOTE_ENGINES = set(_get_remote_engines())

modin/data_management/factories.py

@@ -12,6 +12,8 @@
 # governing permissions and limitations under the License.
 
 import warnings
+import typing
+import re
 
 from modin import execution_engine
 from modin.engines.base.io import BaseIO
@@ -21,13 +23,42 @@
 types_dictionary = {"pandas": {"category": pandas.CategoricalDtype}}
 
 
+class FactoryInfo(typing.NamedTuple):
+    engine: str
+    partition: str
+    experimental: bool
+
+
+class NotRealFactory(Exception):
+    pass
+
+
 class BaseFactory(object):
     """
     Abstract factory which allows to override the io module easily.
     """
 
     io_cls: BaseIO = None  # The module where the I/O functionality exists.
 
+    @classmethod
+    def get_info(cls) -> FactoryInfo:
+        """
+        This gets the information about the factory: its execution engine,
+        partitioning format and whether it's experimental-only.
+
+        Note that it parses factory name, so it must be conformant with how
+        ExecutionEngine class constructs factory names.
+        """
+        try:
+            experimental, partition, engine = re.match(
+                r"^(Experimental)?(.*)On(.*)Factory$", cls.__name__
+            ).groups()
+        except AttributeError:
+            raise NotRealFactory()
+        return FactoryInfo(
+            engine=engine, partition=partition, experimental=bool(experimental)
+        )
+
     @classmethod
     def prepare(cls):
         """
@@ -216,7 +247,9 @@ def prepare(cls):
         cls.io_cls = PyarrowOnRayIO
 
 
-class ExperimentalPandasOnCloudrayFactory(ExperimentalBaseFactory):
+class ExperimentalRemoteFactory(ExperimentalBaseFactory):
+    wrapped_factory = BaseFactory
+
     @classmethod
     def prepare(cls):
         # query_compiler import is needed so remote PandasQueryCompiler
@@ -231,11 +264,13 @@ def prepare(cls):
         import modin.experimental.pandas.numpy_wrap  # noqa: F401
 
         class WrappedIO:
-            def __init__(self, conn):
+            def __init__(self, conn, factory):
                 self.__conn = conn
-                self.__io_cls = conn.modules[
-                    "modin.engines.ray.pandas_on_ray.io"
-                ].PandasOnRayIO
+                remote_factory = getattr(
+                    conn.modules[factory.__module__], factory.__name__
+                )
+                remote_factory.prepare()
+                self.__io_cls = remote_factory.io_cls
                 self.__reads = {
                     name for name in BaseIO.__dict__ if name.startswith("read_")
                 }
@@ -256,4 +291,8 @@ def wrap(*a, _original=getattr(self.__io_cls, name), **kw):
                     wrap = getattr(self.__io_cls, name)
                 return wrap
 
-        cls.io_cls = WrappedIO(get_connection())
+        cls.io_cls = WrappedIO(get_connection(), cls.wrapped_factory)
+
+
+class ExperimentalPandasOnCloudrayFactory(ExperimentalRemoteFactory):
+    wrapped_factory = PandasOnRayFactory

modin/experimental/cloud/meta_magic.py

@@ -16,6 +16,7 @@
 import types
 
 from modin import execution_engine
+from modin.data_management import REMOTE_ENGINES
 
 # the attributes that must be alwasy taken from a local part of dual-nature class,
 # never going to remote end
@@ -153,7 +154,7 @@ def __new__(cls, *a, **kw):
     _KNOWN_DUALS[local_cls] = result
 
     def update_class(_):
-        if execution_engine.get() == "Cloudray":
+        if execution_engine.get() in REMOTE_ENGINES:
             from . import rpyc_proxy
 
             result.__real_cls__ = getattr(rpyc_proxy, rpyc_wrapper_name)(result)

modin/experimental/pandas/numpy_wrap.py

@@ -27,6 +27,7 @@
     import types
     import copyreg
     from modin import execution_engine
+    from modin.data_management import REMOTE_ENGINES
     import modin
     import pandas
     import os
@@ -78,7 +79,7 @@ def __swap_numpy(self, other_numpy=None):
                 self.__has_to_warn = False
 
         def __update_engine(self, _):
-            if execution_engine.get() == "Cloudray":
+            if execution_engine.get() in REMOTE_ENGINES:
                 from modin.experimental.cloud import get_connection
 
                 self.__swap_numpy(get_connection().modules["numpy"])

modin/pandas/__init__.py

@@ -134,9 +134,8 @@ def _update_engine(publisher: Publisher):
 
     elif publisher.get() == "Cloudray":
         from modin.experimental.cloud import get_connection
-        import rpyc
 
-        conn: rpyc.ClassicService = get_connection()
+        conn = get_connection()
         remote_ray = conn.modules["ray"]
         if _is_first_update.get("Cloudray", True):
 
@@ -159,7 +158,6 @@ def init_remote_ray():
             import modin.data_management.dispatcher  # noqa: F401
 
         num_cpus = remote_ray.cluster_resources()["CPU"]
-
     elif publisher.get() not in _NOINIT_ENGINES:
         raise ImportError("Unrecognized execution engine: {}.".format(publisher.get()))