Repository: modin-project/modin

Issue:
Implement a remote function decorator with a cache
Currently, there is no a unified mechanism for putting functions into a distributed store and caching the remote references. There are multiple approaches to solving this issue, for example like here - https://github.com/modin-project/modin/blob/master/modin/core/execution/ray/implementations/pandas_on_ray/partitioning/virtual_partition.py#L64 .

All this functionality could be implemented in a single decorator, for example `@remote_function`, with an internal cache.


Discussion:

Resolving patch:
diff --git a/modin/core/execution/utils.py b/modin/core/execution/utils.py
--- a/modin/core/execution/utils.py
+++ b/modin/core/execution/utils.py
@@ -16,6 +16,8 @@
 import contextlib
 import os
 
+from modin.error_message import ErrorMessage
+
 
 @contextlib.contextmanager
 def set_env(**environ):
@@ -29,3 +31,63 @@ def set_env(**environ):
     finally:
         os.environ.clear()
         os.environ.update(old_environ)
+
+
+if "_MODIN_DOC_CHECKER_" in os.environ:
+
+    # The doc checker should get the non-processed functions
+    def remote_function(func, ignore_defaults=False):
+        return func
+
+
+# Check if the function already exists to avoid circular imports
+elif "remote_function" not in dir():
+    from modin.config import Engine
+
+    if Engine.get() == "Ray":
+        from modin.core.execution.ray.common import RayWrapper
+
+        _preprocess_func = RayWrapper.put
+    elif Engine.get() == "Unidist":
+        from modin.core.execution.unidist.common import UnidistWrapper
+
+        _preprocess_func = UnidistWrapper.put
+    elif Engine.get() == "Dask":
+        from modin.core.execution.dask.common import DaskWrapper
+
+        # The function cache is not supported for Dask
+        def remote_function(func, ignore_defaults=False):
+            return DaskWrapper.put(func)
+
+    else:
+
+        def remote_function(func, ignore_defaults=False):
+            return func
+
+    if "remote_function" not in dir():
+        _remote_function_cache = {}
+
+        def remote_function(func, ignore_defaults=False):  # noqa: F811
+            if "<locals>" in func.__qualname__:  # Nested function
+                if func.__closure__:
+                    ErrorMessage.single_warning(
+                        f"The nested function {func} can not be cached, because "
+                        + "it captures objects from the outer scope."
+                    )
+                    return func
+                if not ignore_defaults and func.__defaults__:
+                    ErrorMessage.single_warning(
+                        f"The nested function {func} can not be cached, because it has "
+                        + "default values. Use `ignore_defaults` to forcibly enable caching."
+                    )
+                    return func
+                # For the nested functions, use __code__ as the key, because it's the same
+                # object for each instance of the function.
+                key = id(func.__code__)
+            else:
+                key = func
+            ref = _remote_function_cache.get(key, None)
+            if ref is None:
+                ref = _preprocess_func(func)
+                _remote_function_cache[key] = ref
+            return ref
diff --git a/scripts/doc_checker.py b/scripts/doc_checker.py
--- a/scripts/doc_checker.py
+++ b/scripts/doc_checker.py
@@ -35,6 +35,9 @@
 from numpydoc.docscrape import NumpyDocString
 from numpydoc.validate import Docstring
 
+# Let the other modules to know that the doc checker is running.
+os.environ["_MODIN_DOC_CHECKER_"] = "1"
+
 # fake cuDF-related modules if they're missing
 for mod_name in ("cudf", "cupy"):
     try:


Test patch:
diff --git a/modin/test/storage_formats/pandas/test_internals.py b/modin/test/storage_formats/pandas/test_internals.py
--- a/modin/test/storage_formats/pandas/test_internals.py
+++ b/modin/test/storage_formats/pandas/test_internals.py
@@ -28,6 +28,7 @@
     LazyProxyCategoricalDtype,
     ModinDtypes,
 )
+from modin.core.execution.utils import remote_function
 from modin.core.storage_formats.pandas.utils import split_result_of_axis_func_pandas
 from modin.distributed.dataframe.pandas import from_partitions
 from modin.pandas.test.utils import (
@@ -55,6 +56,8 @@
     virtual_column_partition_class = PandasOnRayDataframeColumnPartition
     virtual_row_partition_class = PandasOnRayDataframeRowPartition
     put = RayWrapper.put
+    deploy = RayWrapper.deploy
+    materialize = RayWrapper.materialize
 elif Engine.get() == "Dask":
     from modin.core.execution.dask.common import DaskWrapper
     from modin.core.execution.dask.implementations.pandas_on_dask.partitioning import (
@@ -72,6 +75,8 @@ def put(x):
     block_partition_class = PandasOnDaskDataframePartition
     virtual_column_partition_class = PandasOnDaskDataframeColumnPartition
     virtual_row_partition_class = PandasOnDaskDataframeRowPartition
+    deploy = DaskWrapper.deploy
+    materialize = DaskWrapper.materialize
 elif Engine.get() == "Python":
     from modin.core.execution.python.common import PythonWrapper
     from modin.core.execution.python.implementations.pandas_on_python.partitioning import (
@@ -83,6 +88,12 @@ def put(x):
     def put(x):
         return PythonWrapper.put(x, hash=False)
 
+    def deploy(func, args=tuple()):
+        return func(*args)
+
+    def materialize(arg):
+        return arg
+
     block_partition_class = PandasOnPythonDataframePartition
     virtual_column_partition_class = PandasOnPythonDataframeColumnPartition
     virtual_row_partition_class = PandasOnPythonDataframeRowPartition
@@ -2507,3 +2518,36 @@ def f2(i):
 
     hook = MetaList(f1.remote())[2]
     assert ray.get(f2.remote(hook)) == 3
+
+
+def test_remote_function():
+    def get_func():
+        @remote_function
+        def remote_func(arg):
+            return arg
+
+        return remote_func
+
+    def get_capturing_func(arg):
+        @remote_function
+        def remote_func():
+            return arg
+
+        return remote_func
+
+    if Engine.get() in ("Ray", "Unidist"):
+        from modin.core.execution.utils import _remote_function_cache
+
+        cache_len = len(_remote_function_cache)
+        assert get_func() is get_func()
+        assert get_func() in _remote_function_cache.values()
+        assert get_capturing_func(1) not in _remote_function_cache.values()
+        assert len(_remote_function_cache) == cache_len + 1
+
+    assert materialize(deploy(get_func(), [123])) == 123
+    assert get_capturing_func(1) is not get_capturing_func(2)
+    assert (
+        materialize(deploy(get_capturing_func(1)))
+        + materialize(deploy(get_capturing_func(2)))
+        == 3
+    )