Implement post task gc hook (#47)

* Implement post task gc hook

* Address Comments

* Move gc import to the top

modin/pandas/groupby.py

@@ -12,7 +12,7 @@
 
 from .concat import concat
 from .index_metadata import _IndexMetadata
-from .utils import _inherit_docstrings, _reindex_helper
+from .utils import _inherit_docstrings, _reindex_helper, post_task_gc
 
 
 @_inherit_docstrings(pandas.core.groupby.DataFrameGroupBy,
@@ -553,6 +553,7 @@ def _apply_df_function(self, f, concat_axis=None):
 
 
 @ray.remote
+@post_task_gc
 def groupby(by, axis, level, as_index, sort, group_keys, squeeze, *df):
 
     df = pandas.concat(df, axis=axis)

modin/pandas/utils.py

@@ -7,6 +7,8 @@
 import collections
 import numpy as np
 import ray
+import time
+import gc
 
 from . import get_npartitions
 
@@ -104,6 +106,37 @@ def remote(self, *args):
         return result
 
 
+def post_task_gc(func):
+    """Perform garbage collection after the task is executed.
+
+    Usage:
+        ```
+        @ray.remote
+        @post_task_gc
+        def memory_hungry_op():
+            ...
+        ```
+    Note:
+        - This will invoke the GC for the entire process. Expect
+          About 100ms latency.
+        - We have a basic herustic in place to balance of tradeoff between
+          speed and memory. If the task takes more than 500ms to run, we
+          will do the GC.
+    """
+    def wrapped(*args):
+        start_time = time.time()
+
+        result = func(*args)
+
+        duration_s = time.time() - start_time
+        duration_ms = duration_s * 1000
+        if duration_ms > 500:
+            gc.collect()
+
+        return result
+    return wrapped
+
+
 def _get_nan_block_id(n_row=1, n_col=1, transpose=False):
     """A memory efficent way to get a block of NaNs.