Repository: mars-project/mars

Issue:
[BUG] xgboost on Mars will try fetching all DMatrices onto tracker machine
**Describe the bug**
xgboost on Mars will try fetching all DMatrices onto tracker machine, making the whole task stuck.


Discussion:

Resolving patch:
diff --git a/mars/services/scheduling/worker/execution.py b/mars/services/scheduling/worker/execution.py
--- a/mars/services/scheduling/worker/execution.py
+++ b/mars/services/scheduling/worker/execution.py
@@ -167,15 +167,10 @@ async def _prepare_input_data(self, subtask: Subtask, band_name: str):
         storage_api = await StorageAPI.create(
             subtask.session_id, address=self.address, band_name=band_name
         )
-        pure_dep_keys = set()
         chunk_key_to_data_keys = get_chunk_key_to_data_keys(subtask.chunk_graph)
-        for n in subtask.chunk_graph:
-            pure_dep_keys.update(
-                inp.key
-                for inp, pure_dep in zip(n.inputs, n.op.pure_depends)
-                if pure_dep
-            )
         for chunk in subtask.chunk_graph:
+            if chunk.key in subtask.pure_depend_keys:
+                continue
             if chunk.op.gpu:  # pragma: no cover
                 to_fetch_band = band_name
             else:
@@ -205,7 +200,11 @@ async def _collect_input_sizes(
         sizes = dict()
 
         fetch_keys = list(
-            set(n.key for n in graph.iter_indep() if isinstance(n.op, Fetch))
+            set(
+                n.key
+                for n in graph.iter_indep()
+                if isinstance(n.op, Fetch) and n.key not in subtask.pure_depend_keys
+            )
         )
         if not fetch_keys:
             return sizes
@@ -251,6 +250,8 @@ def _estimate_sizes(cls, subtask: Subtask, input_sizes: Dict):
         # condense op key graph
         op_key_graph = DAG()
         for n in graph.topological_iter():
+            if n.key in subtask.pure_depend_keys:
+                continue
             if n.op.key not in op_key_graph:
                 op_key_graph.add_node(n.op.key)
             for succ in graph.iter_successors(n):
diff --git a/mars/services/subtask/core.py b/mars/services/subtask/core.py
--- a/mars/services/subtask/core.py
+++ b/mars/services/subtask/core.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from enum import Enum
-from typing import Iterable, List, Optional, Tuple
+from typing import Iterable, List, Optional, Set, Tuple
 
 from ...core import ChunkGraph, DAG
 from ...serialization.serializables import (
@@ -50,7 +50,7 @@ def is_done(self) -> bool:
 
 
 class Subtask(Serializable):
-    __slots__ = ("_repr",)
+    __slots__ = ("_repr", "_pure_depend_keys")
 
     subtask_id: str = StringField("subtask_id")
     subtask_name: str = StringField("subtask_name")
@@ -110,6 +110,7 @@ def __init__(
             logic_parallelism=logic_parallelism,
             bands_specified=bands_specified,
         )
+        self._pure_depend_keys = None
         self._repr = None
 
     @property
@@ -117,6 +118,20 @@ def expect_band(self):
         if self.expect_bands:
             return self.expect_bands[0]
 
+    @property
+    def pure_depend_keys(self) -> Set[str]:
+        if self._pure_depend_keys is not None:
+            return self._pure_depend_keys
+        pure_dep_keys = set()
+        for n in self.chunk_graph:
+            pure_dep_keys.update(
+                inp.key
+                for inp, pure_dep in zip(n.inputs, n.op.pure_depends)
+                if pure_dep
+            )
+        self._pure_depend_keys = pure_dep_keys
+        return pure_dep_keys
+
     def __repr__(self):
         if self._repr is not None:
             return self._repr
@@ -145,7 +160,7 @@ class SubtaskResult(Serializable):
     bands: List[BandType] = ListField("band", FieldTypes.tuple, default=None)
     error = AnyField("error", default=None)
     traceback = AnyField("traceback", default=None)
-    # The following is the execution information of the sub task
+    # The following is the execution information of the subtask
     execution_start_time: float = Float64Field("execution_start_time")
     execution_end_time: float = Float64Field("execution_end_time")
 
diff --git a/mars/services/subtask/worker/processor.py b/mars/services/subtask/worker/processor.py
--- a/mars/services/subtask/worker/processor.py
+++ b/mars/services/subtask/worker/processor.py
@@ -121,7 +121,10 @@ def subtask_id(self):
     async def _load_input_data(self):
         keys, gets, accept_nones = [], [], []
         for chunk in self._chunk_graph.iter_indep():
-            if isinstance(chunk.op, Fetch):
+            if (
+                isinstance(chunk.op, Fetch)
+                and chunk.key not in self.subtask.pure_depend_keys
+            ):
                 keys.append(chunk.key)
                 gets.append(self._storage_api.get.delay(chunk.key))
                 accept_nones.append(True)


Test patch:
diff --git a/mars/services/scheduling/worker/tests/test_execution.py b/mars/services/scheduling/worker/tests/test_execution.py
--- a/mars/services/scheduling/worker/tests/test_execution.py
+++ b/mars/services/scheduling/worker/tests/test_execution.py
@@ -391,6 +391,35 @@ def delay_fun(delay, _inp1):
     )
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("actor_pool", [(1, True)], indirect=True)
+async def test_execute_with_pure_deps(actor_pool):
+    pool, session_id, meta_api, storage_api, execution_ref = actor_pool
+
+    dep = TensorFetch(key="input1", dtype=np.dtype(int)).new_chunk([])
+
+    def main_fun():
+        return session_id
+
+    remote_result = RemoteFunction(
+        function=main_fun, function_args=[], function_kwargs={}
+    ).new_chunk([dep])
+    # mark `dep` as pure dependency
+    remote_result.op._pure_depends = [True]
+    chunk_graph = ChunkGraph([remote_result])
+    chunk_graph.add_node(dep)
+    chunk_graph.add_node(remote_result)
+    chunk_graph.add_edge(dep, remote_result)
+
+    subtask = Subtask(
+        f"test_subtask_{uuid.uuid4()}", session_id=session_id, chunk_graph=chunk_graph
+    )
+    # subtask shall run well without data of `dep` available
+    await execution_ref.run_subtask(subtask, "numa-0", pool.external_address)
+    res = await storage_api.get(remote_result.key)
+    assert res == session_id
+
+
 def test_estimate_size():
     from ..execution import SubtaskExecutionActor
     from .....dataframe.arithmetic import DataFrameAdd