Repository: mars-project/mars

Issue:
[BUG] Memory size estimation for TensorDot is missing
**Describe the bug**
When merging executables, function ``_tensordot_estimate_size`` is missing which makes computation graphs with dense ``TensorDot``s easy to crash. We need to add this estimation back.


Discussion:

Resolving patch:
diff --git a/mars/executor.py b/mars/executor.py
--- a/mars/executor.py
+++ b/mars/executor.py
@@ -439,7 +439,7 @@ def _preprocess(self, graph, keys):
 
     @classmethod
     def handle(cls, op, results, mock=False):
-        method_name, mapper = ('execute', cls._op_runners) if not mock else\
+        method_name, mapper = ('execute', cls._op_runners) if not mock else \
             ('estimate_size', cls._op_size_estimators)
         try:
             runner = mapper[type(op)]
@@ -642,13 +642,19 @@ def ignore(*_):
 
 
 Executor._op_runners[Fetch] = ignore
-Executor._op_runners[ShuffleProxy] =ignore
+Executor._op_runners[ShuffleProxy] = ignore
 
 
-def register(op, handler, size_estimator=None):
-    Executor._op_runners[op] = handler
+def register(op_cls, handler=None, size_estimator=None):
+    if handler:
+        Executor._op_runners[op_cls] = handler
     if size_estimator:
-        Executor._op_size_estimators[op] = size_estimator
+        Executor._op_size_estimators[op_cls] = size_estimator
+
+
+def register_default(op_cls):
+    Executor._op_runners.pop(op_cls, None)
+    Executor._op_size_estimators.pop(op_cls, None)
 
 
 # import to register operands
diff --git a/mars/tensor/linalg/tensordot.py b/mars/tensor/linalg/tensordot.py
--- a/mars/tensor/linalg/tensordot.py
+++ b/mars/tensor/linalg/tensordot.py
@@ -68,6 +68,22 @@ def __call__(self, a, b):
             tuple(s for i, s in enumerate(b.shape) if i not in set(self._b_axes))
         return self.new_tensor([a, b], shape, order=TensorOrder.C_ORDER)
 
+    @classmethod
+    def estimate_size(cls, ctx, op):
+        chunk = op.outputs[0]
+        if chunk.is_sparse():
+            return super(TensorTensorDot, cls).estimate_size(ctx, op)
+
+        # empirical value in real environments
+        calc_usage = chunk.nbytes
+
+        # add input sizes when sparse-to-dense is needed
+        for inp in chunk.inputs:
+            if inp.is_sparse():
+                calc_usage += inp.nbytes
+
+        ctx[chunk.key] = (chunk.nbytes, calc_usage)
+
     @classmethod
     def tile(cls, op):
         a, b, a_axes, b_axes = op.a, op.b, op.a_axes, op.b_axes


Test patch:
diff --git a/mars/tensor/linalg/tests/test_linalg_execute.py b/mars/tensor/linalg/tests/test_linalg_execute.py
--- a/mars/tensor/linalg/tests/test_linalg_execute.py
+++ b/mars/tensor/linalg/tests/test_linalg_execute.py
@@ -820,9 +820,11 @@ def testTensordotExecution(self):
         axes = ([1, 0], [0, 1])
         c = tensordot(a, b, axes=axes)
         size_res = size_executor.execute_tensor(c, mock=True)
+        self.assertEqual(sum(s[0] for s in size_res), c.nbytes)
+        self.assertEqual(sum(s[1] for s in size_res), c.nbytes)
+
         res = self.executor.execute_tensor(c)
         expected = np.tensordot(a_data, b_data, axes=axes)
-        self.assertEqual(sum(s[0] for s in size_res), c.nbytes)
         self.assertTrue(np.array_equal(res[0], expected[:2, :]))
         self.assertTrue(np.array_equal(res[1], expected[2:4, :]))
         self.assertTrue(np.array_equal(res[2], expected[4:, :]))
@@ -874,28 +876,71 @@ def testTensordotExecution(self):
         res = self.executor.execute_tensor(c, concat=True)[0]
         np.testing.assert_array_equal(res, np.ones((100, 100)) * 100)
 
-    def testSparseDotExecution(self):
-        size_executor = Executor(sync_provider_type=Executor.SyncProviderType.MOCK)
+    def testSparseDotSizeExecution(self):
+        from mars.tensor.linalg.tensordot import TensorTensorDot
+        from mars.executor import register, register_default
+        chunk_sizes = dict()
+        chunk_nbytes = dict()
+        chunk_input_sizes = dict()
+        chunk_input_nbytes = dict()
+
+        def execute_size(t):
+            def _tensordot_size_recorder(ctx, op):
+                TensorTensorDot.estimate_size(ctx, op)
+
+                chunk_key = op.outputs[0].key
+                chunk_sizes[chunk_key] = ctx[chunk_key]
+                chunk_nbytes[chunk_key] = op.outputs[0].nbytes
+
+                input_sizes = dict((inp.op.key, ctx[inp.key][0]) for inp in op.inputs)
+                chunk_input_sizes[chunk_key] = sum(input_sizes.values())
+                input_nbytes = dict((inp.op.key, inp.nbytes) for inp in op.inputs)
+                chunk_input_nbytes[chunk_key] = sum(input_nbytes.values())
+
+            size_executor = Executor(sync_provider_type=Executor.SyncProviderType.MOCK)
+            try:
+                chunk_sizes.clear()
+                chunk_nbytes.clear()
+                chunk_input_sizes.clear()
+                chunk_input_nbytes.clear()
+                register(TensorTensorDot, size_estimator=_tensordot_size_recorder)
+                size_executor.execute_tensor(t, mock=True)
+            finally:
+                register_default(TensorTensorDot)
 
         a_data = sps.random(5, 9, density=.1)
         b_data = sps.random(9, 10, density=.2)
         a = tensor(a_data, chunk_size=2)
         b = tensor(b_data, chunk_size=3)
 
+        c = dot(a, b)
+        execute_size(c)
+
+        for key in chunk_input_sizes.keys():
+            self.assertGreaterEqual(chunk_sizes[key][1], chunk_input_sizes[key])
+
+        c2 = dot(a, b, sparse=False)
+        execute_size(c2)
+
+        for key in chunk_input_sizes.keys():
+            self.assertEqual(chunk_sizes[key][0], chunk_nbytes[key])
+            self.assertEqual(chunk_sizes[key][1], chunk_input_nbytes[key] + chunk_nbytes[key])
+
+    def testSparseDotExecution(self):
+        a_data = sps.random(5, 9, density=.1)
+        b_data = sps.random(9, 10, density=.2)
+        a = tensor(a_data, chunk_size=2)
+        b = tensor(b_data, chunk_size=3)
+
         c = dot(a, b)
 
-        size_res = size_executor.execute_tensor(c, mock=True)
         res = self.executor.execute_tensor(c, concat=True)[0]
-        self.assertEqual(sum(s[0] for s in size_res), 0)
-        self.assertGreaterEqual(sum(s[1] for s in size_res), 0)
         self.assertTrue(issparse(res))
         np.testing.assert_allclose(res.toarray(), a_data.dot(b_data).toarray())
 
         c2 = dot(a, b, sparse=False)
 
-        size_res = size_executor.execute_tensor(c2, mock=True)
         res = self.executor.execute_tensor(c2, concat=True)[0]
-        self.assertEqual(sum(s[0] for s in size_res), c2.nbytes)
         self.assertFalse(issparse(res))
         np.testing.assert_allclose(res, a_data.dot(b_data).toarray())