Repository: mars-project/mars

Issue:
[BUG] Failed to execute one of the multiple outputs
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**
Operator like `mt.linalg.qr` has multiple outputs, if execute one of the output tensors, it will raise a `KeyError`.

**To Reproduce**
``` Python
In [1]: import mars.tensor as mt                                                

In [2]: a = mt.random.randint(1, 10, (5, 5)) 
   ...: arrs = mt.linalg.qr(a) 
   ...: arrs[0].execute()                                                       
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
~/Documents/mars_dev/mars/mars/graph.pyx in mars.graph.DirectedGraph.iter_predecessors()

KeyError: Chunk <op=TensorQR, key=fd777e263eae8d37f1d6057a36de12a4>

During handling of the above exception, another exception occurred:

KeyError                                  Traceback (most recent call last)
<ipython-input-2-17af77a63ed7> in <module>
      1 a = mt.random.randint(1, 10, (5, 5))
      2 arrs = mt.linalg.qr(a)
----> 3 arrs[0].execute()

~/Documents/mars_dev/mars/mars/tensor/core.py in execute(self, session, **kw)
    275         if session is None:
    276             session = Session.default_or_local()
--> 277         return session.run(self, **kw)
    278 
    279     def fetch(self, session=None, **kw):

~/Documents/mars_dev/mars/mars/session.py in run(self, *tensors, **kw)
    107 
    108         tensors = tuple(mt.tensor(t) for t in tensors)
--> 109         result = self._sess.run(*tensors, **kw)
    110 
    111         for t in tensors:

~/Documents/mars_dev/mars/mars/session.py in run(self, *tensors, **kw)
     49         if 'n_parallel' not in kw:
     50             kw['n_parallel'] = cpu_count()
---> 51         res = self._executor.execute_tensors(tensors, **kw)
     52         return res
     53 

~/Documents/mars_dev/mars/mars/utils.py in _wrapped(*args, **kwargs)
    343     def _wrapped(*args, **kwargs):
    344         _kernel_mode.eager = False
--> 345         return_value = func(*args, **kwargs)
    346         _kernel_mode.eager = None
    347         return return_value

~/Documents/mars_dev/mars/mars/executor.py in execute_tensors(self, tensors, fetch, n_parallel, n_thread, print_progress, mock, sparse_mock_percent)
    470         self.execute_graph(graph, result_keys, n_parallel=n_parallel or n_thread,
    471                            print_progress=print_progress, mock=mock,
--> 472                            sparse_mock_percent=sparse_mock_percent)
    473 
    474         results = self._chunk_result

~/Documents/mars_dev/mars/mars/executor.py in execute_graph(self, graph, keys, n_parallel, print_progress, mock, sparse_mock_percent)
    407                                          print_progress=print_progress, mock=mock,
    408                                          sparse_mock_percent=sparse_mock_percent)
--> 409         return graph_execution.execute(True)
    410 
    411     @kernel_mode

~/Documents/mars_dev/mars/mars/executor.py in execute(self, retval)
    336         # wait until all the futures completed
    337         for future in executed_futures:
--> 338             future.result()
    339 
    340         if self._mock:

~/miniconda3/lib/python3.7/concurrent/futures/_base.py in result(self, timeout)
    423                 raise CancelledError()
    424             elif self._state == FINISHED:
--> 425                 return self.__get_result()
    426 
    427             self._condition.wait(timeout)

~/miniconda3/lib/python3.7/concurrent/futures/_base.py in __get_result(self)
    382     def __get_result(self):
    383         if self._exception:
--> 384             raise self._exception
    385         else:
    386             return self._result

~/miniconda3/lib/python3.7/concurrent/futures/thread.py in run(self)
     55 
     56         try:
---> 57             result = self.fn(*self.args, **self.kwargs)
     58         except BaseException as exc:
     59             self.future.set_exception(exc)

~/Documents/mars_dev/mars/mars/executor.py in _execute_operand(self, op)
    246 
    247                     # clean the predecessors' results if ref counts equals 0
--> 248                     for pred_chunk in self._graph.iter_predecessors(output):
    249                         if pred_chunk.key in ref_counts:
    250                             ref_counts[pred_chunk.key] -= 1

~/Documents/mars_dev/mars/mars/graph.pyx in mars.graph.DirectedGraph.iter_predecessors()

KeyError: 'Node Chunk <op=TensorQR, key=fd777e263eae8d37f1d6057a36de12a4> does not exist in the directed graph'
```



Discussion:

Resolving patch:
diff --git a/mars/executor.py b/mars/executor.py
--- a/mars/executor.py
+++ b/mars/executor.py
@@ -239,6 +239,9 @@ def _execute_operand(self, op):
 
             with self._lock:
                 for output in itertools.chain(*[op.outputs for op in ops]):
+                    # the output not in the graph will be skipped
+                    if output not in self._graph:
+                        continue
                     # in case that operand has multiple outputs
                     # and some of the output not in result keys, delete them
                     if ref_counts.get(output.key) == 0:


Test patch:
diff --git a/mars/tests/test_session.py b/mars/tests/test_session.py
--- a/mars/tests/test_session.py
+++ b/mars/tests/test_session.py
@@ -55,6 +55,15 @@ def testMultipleOutputExecute(self):
 
         np.testing.assert_array_equal(result, expected)
 
+        # test multiple outputs, but only execute 1
+        data = np.random.randint(0, 10, (5, 5))
+        arr3 = mt.tensor(data)
+        arrs = mt.linalg.qr(arr3)
+        result = arrs[0].execute()
+        expected = np.linalg.qr(data)[0]
+
+        np.testing.assert_array_almost_equal(result, expected)
+
     def testReExecuteSame(self):
         data = np.random.random((5, 9))