Repository: mars-project/mars

Issue:
[BUG] StartTracker does not specify `pure_depends` which leads to wrong memory estimation
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**

`StartTracker` in XGB is actually a virtual node without memory consumption, and it depends on all the data, however, `pure_depends` is not specified so that the memory estimation is equal to all the input data size.

**To Reproduce**
To help us reproducing this bug, please provide information below:
1. Your Python version
2. The version of Mars you use
3. Versions of crucial packages, such as numpy, scipy and protobuf
4. Full stack of the error.
5. Minimized code to reproduce the error.

```
2021-01-26 15:39:43,801 mars.scheduler.operands.base 474 DEBUG    Operand e3036e9fb2b4f84911306f8fa9cf3e03(StartTracker) state from OperandState.UNSCHEDULED to OperandState.READY.
2021-01-26 15:39:43,802 mars.scheduler.assigner 476 DEBUG    Operand e3036e9fb2b4f84911306f8fa9cf3e03 enqueued

2021-01-26 15:49:43,403 mars.scheduler.assigner 473 ERROR    Unexpected exception occurred in AssignEvaluationActor._allocate_resource. op_key=e3036e9fb2b4f84911306f8fa9cf3e03
Traceback (most recent call last):
  File "/home/admin/work/_public-mars-0.6.2.zip/mars/utils.py", line 377, in _wrapped
    return func(*args, **kwargs)
  File "/home/admin/work/_public-mars-0.6.2.zip/mars/scheduler/assigner.py", line 414, in _allocate_resource
    raise TimeoutError(f'Assign resources to operand {op_key} timed out')
TimeoutError: Assign resources to operand e3036e9fb2b4f84911306f8fa9cf3e03 timed out
2021-01-26 15:49:43,404 mars.scheduler.assigner 473 ERROR    Unexpected error occurred in s:0:AssignEvaluationActor$eb2a35c62f8cfb405cc15245e0d1cf02
Traceback (most recent call last):
  File "/home/admin/work/_public-mars-0.6.2.zip/mars/scheduler/assigner.py", line 315, in allocate_top_resources
    reject_workers=reject_workers)
  File "/home/admin/work/_public-mars-0.6.2.zip/mars/utils.py", line 377, in _wrapped
    return func(*args, **kwargs)
  File "/home/admin/work/_public-mars-0.6.2.zip/mars/scheduler/assigner.py", line 414, in _allocate_resource
    raise TimeoutError(f'Assign resources to operand {op_key} timed out')
TimeoutError: Assign resources to operand e3036e9fb2b4f84911306f8fa9cf3e03 timed out
2021-01-26 15:49:43,407 mars.scheduler.operands.common 474 ERROR    Attempt 1: Unexpected error TimeoutError occurred in executing operand e3036e9fb2b4f84911306f8fa9cf3e03 in None
Traceback (most recent call last):
  File "/home/admin/work/_public-mars-0.6.2.zip/mars/scheduler/assigner.py", line 315, in allocate_top_resources
    reject_workers=reject_workers)
  File "/home/admin/work/_public-mars-0.6.2.zip/mars/utils.py", line 377, in _wrapped
    return func(*args, **kwargs)
  File "/home/admin/work/_public-mars-0.6.2.zip/mars/scheduler/assigner.py", line 414, in _allocate_resource
    raise TimeoutError(f'Assign resources to operand {op_key} timed out')
TimeoutError: Assign resources to operand e3036e9fb2b4f84911306f8fa9cf3e03 timed out
```


Discussion:

Resolving patch:
diff --git a/mars/learn/contrib/xgboost/start_tracker.py b/mars/learn/contrib/xgboost/start_tracker.py
--- a/mars/learn/contrib/xgboost/start_tracker.py
+++ b/mars/learn/contrib/xgboost/start_tracker.py
@@ -27,8 +27,9 @@ class StartTracker(LearnOperand, LearnOperandMixin):
 
     _n_workers = Int32Field('n_workers')
 
-    def __init__(self, n_workers=None, output_types=None, **kw):
-        super().__init__(_n_workers=n_workers, _output_types=output_types, **kw)
+    def __init__(self, n_workers=None, output_types=None, pure_depends=None, **kw):
+        super().__init__(_n_workers=n_workers, _output_types=output_types,
+                         _pure_depends=pure_depends, **kw)
         if self.output_types is None:
             self.output_types = [OutputType.object]
 
diff --git a/mars/learn/contrib/xgboost/train.py b/mars/learn/contrib/xgboost/train.py
--- a/mars/learn/contrib/xgboost/train.py
+++ b/mars/learn/contrib/xgboost/train.py
@@ -118,7 +118,9 @@ def tile(cls, op):
             inp = op.inputs[0]
             in_chunks = inp.chunks
             workers = cls._get_dmatrix_chunks_workers(ctx, inp)
-            tracker_chunk = StartTracker(n_workers=len(in_chunks)).new_chunk(in_chunks, shape=())
+            n_chunk = len(in_chunks)
+            tracker_chunk = StartTracker(n_workers=n_chunk, pure_depends=[True] * n_chunk)\
+                .new_chunk(in_chunks, shape=())
             out_chunks = []
             worker_to_evals = defaultdict(list)
             if op.evals is not None:


Test patch:
diff --git a/mars/learn/contrib/xgboost/tests/integrated/test_distributed_xgboost.py b/mars/learn/contrib/xgboost/tests/integrated/test_distributed_xgboost.py
--- a/mars/learn/contrib/xgboost/tests/integrated/test_distributed_xgboost.py
+++ b/mars/learn/contrib/xgboost/tests/integrated/test_distributed_xgboost.py
@@ -19,16 +19,29 @@
 
 import mars.dataframe as md
 import mars.tensor as mt
+from mars.executor import register
 from mars.session import new_session
 from mars.learn.contrib.xgboost import XGBClassifier
 from mars.tests.integrated.base import IntegrationTestBase
+from mars.utils import calc_data_size
 
 try:
     import xgboost
+
+    from mars.learn.contrib.xgboost.start_tracker import StartTracker
 except ImportError:
     xgboost = None
 
 
+if xgboost and os.environ.get('TEST_START_TRACKER') == '1':
+    def _patch_start_tracker_estimator(ctx, op: StartTracker):
+        op.estimate_size(ctx, op)
+        estimated_size = ctx[op.outputs[0].key]
+        assert estimated_size[0] == estimated_size[1] == calc_data_size(op.outputs[0])
+
+    register(StartTracker, StartTracker.execute, _patch_start_tracker_estimator)
+
+
 @unittest.skipIf(xgboost is None, 'xgboost not installed')
 class Test(IntegrationTestBase):
     def setUp(self):
@@ -40,6 +53,18 @@ def setUp(self):
         self.y = rs.rand(n_rows, chunk_size=chunk_size)
         super().setUp()
 
+    @property
+    def _extra_worker_options(self):
+        return ['--load-modules',
+                'mars.learn.contrib.xgboost.tests.'
+                'integrated.test_distributed_xgboost']
+
+    @property
+    def _worker_env(self):
+        env = os.environ.copy()
+        env['TEST_START_TRACKER'] = '1'
+        return env
+
     def testDistributedXGBClassifier(self):
         service_ep = 'http://127.0.0.1:' + self.web_port
         timeout = 120 if 'CI' in os.environ else -1
diff --git a/mars/tests/integrated/base.py b/mars/tests/integrated/base.py
--- a/mars/tests/integrated/base.py
+++ b/mars/tests/integrated/base.py
@@ -63,6 +63,14 @@ def _extra_scheduler_options(self):
     def _extra_worker_options(self):
         return []
 
+    @property
+    def _scheduler_env(self):
+        return
+
+    @property
+    def _worker_env(self):
+        return
+
     def _start_distributed_env(self, n_workers=2):
         scheduler_port = self.scheduler_port = str(get_next_port())
         self.proc_workers = []
@@ -76,7 +84,8 @@ def _start_distributed_env(self, n_workers=2):
                                             '--schedulers', '127.0.0.1:' + scheduler_port,
                                             '--log-level', 'debug',
                                             '--log-format', f'WOR{idx} %(asctime)-15s %(message)s',
-                                            '--ignore-avail-mem'] + self._extra_worker_options)
+                                            '--ignore-avail-mem'] + self._extra_worker_options,
+                                           env=self._worker_env)
 
             self.proc_workers.append(proc_worker)
 
@@ -88,7 +97,8 @@ def _start_distributed_env(self, n_workers=2):
                                            '-Dscheduler.retry_delay=5',
                                            '--log-level', 'debug',
                                            '--log-format', 'SCH %(asctime)-15s %(message)s'] +
-                                          self._extra_scheduler_options)
+                                          self._extra_scheduler_options,
+                                          env=self._scheduler_env)
         self.proc_scheduler = proc_scheduler
 
         self.wait_scheduler_worker_start()