Repository: mars-project/mars

Issue:
[BUG][DISTRIBUTED] run parts of svd results will cause serialize error
**Describe the bug**

Some tensor op can generate multiple outputs, svd etc. Assume u, s, v is the result tensors for SVD decomposition, for the distributed setting, if we only run parts of the tensors, for example, u and s, the KeyError will be raised which I guess is all due to the serialization issue.

**To Reproduce**

I start one scheduler, worker and web on my laptop.

On the client side:

```python
In [3]: from mars.session import new_session

In [4]: session  = new_session('http://0.0.0.0:61378')

In [5]: a = mt.random.rand(10, 20)

In [6]: dec = np.linalg.svd(a)

In [7]: dec = mt.linalg.svd(a)

In [8]: session.run(dec[:2])
```

The client will hang.

On the scheduler side, we can get the error:

```python
Traceback (most recent call last):
  File "src/gevent/greenlet.py", line 716, in gevent._greenlet.Greenlet.run
  File "mars/actors/pool/gevent_pool.pyx", line 88, in mars.actors.pool.gevent_pool.ActorExecutionContext.fire_run
    cpdef object fire_run(self):
  File "mars/actors/pool/gevent_pool.pyx", line 91, in mars.actors.pool.gevent_pool.ActorExecutionContext.fire_run
    with self.lock:
  File "mars/actors/pool/gevent_pool.pyx", line 102, in mars.actors.pool.gevent_pool.ActorExecutionContext.fire_run
    raise
  File "mars/actors/pool/gevent_pool.pyx", line 96, in mars.actors.pool.gevent_pool.ActorExecutionContext.fire_run
    res = actor.on_receive(message_ctx.message)
  File "mars/actors/core.pyx", line 108, in mars.actors.core.FunctionActor.on_receive
    cpdef on_receive(self, message):
  File "mars/actors/core.pyx", line 110, in mars.actors.core.FunctionActor.on_receive
    return getattr(self, method)(*args, **kwargs)
  File "/Users/qinxuye/Workspace/mars/mars/scheduler/graph.py", line 153, in execute_graph
    self.prepare_graph()
  File "/Users/qinxuye/Workspace/mars/mars/scheduler/graph.py", line 207, in prepare_graph
    tensor_graph = deserialize_graph(self._serialized_tensor_graph)
  File "/Users/qinxuye/Workspace/mars/mars/utils.py", line 403, in deserialize_graph
    return DirectedGraph.from_json(json_obj)
  File "mars/graph.pyx", line 415, in mars.graph.DirectedGraph.from_json
    return cls.deserialize(SerialiableGraph.from_json(json_obj))
  File "mars/serialize/core.pyx", line 537, in mars.serialize.core.Serializable.from_json
    return cls.deserialize(provider, obj)
  File "mars/serialize/core.pyx", line 510, in mars.serialize.core.Serializable.deserialize
    [cb(key_to_instance) for cb in callbacks]
  File "mars/serialize/jsonserializer.pyx", line 176, in mars.serialize.jsonserializer.JsonSerializeProvider._deserialize_list.cb.inner
    o = subs[v.key, v.id]
KeyError: ('9e2478de2695b727435601490ebaa999', '5219323976')
2018-12-07T11:14:35Z <Greenlet "Greenlet-0" at 0x138f5a448: <built-in method fire_run of mars.actors.pool.gevent_pool.ActorExecutionContext object at 0x139002cc8>> failed with KeyError
```





Discussion:
The failure can be produced without starting the distributed runtime by the code below.

```python
a = mt.random.rand(9, 6, chunks=(3, 6))
U, s, V = mt.linalg.svd(a)

graph = DirectedGraph()
U.build_graph(tiled=False, graph=graph)
s.build_graph(tiled=False, graph=graph)

new_graph = DirectedGraph.from_json(graph.to_json())
```

I will add it to the expression tests and try to fix it.
This issue is not solved completely, for example

```python
a = mt.random.rand(9, 6, chunks=(3, 6))
U, s, V = mt.linalg.svd(a)

graph = DirectedGraph()
(U + 1).build_graph(tiled=False, graph=graph)
(s + 1).build_graph(tiled=False, graph=graph)

new_graph = DirectedGraph.from_json(graph.to_json())
```

We do some calculation on the U and s, the serialize error will throw again, I will try to fix it.

Resolving patch:
diff --git a/mars/tensor/core.py b/mars/tensor/core.py
--- a/mars/tensor/core.py
+++ b/mars/tensor/core.py
@@ -303,7 +303,8 @@ def build_graph(self, graph=None, cls=DAG, tiled=False, compose=True):
                     graph.add_node(c)
                 if not graph.has_successor(c, chunk):
                     graph.add_edge(c, chunk)
-            nodes.extend([c for c in children if c not in visited])
+            nodes.extend([c for c in itertools.chain(*[inp.op.outputs for inp in chunk.inputs or []])
+                          if c not in visited])
         if tiled and compose:
             graph.compose(keys=keys)
         return graph


Test patch:
diff --git a/mars/deploy/local/tests/test_cluster.py b/mars/deploy/local/tests/test_cluster.py
--- a/mars/deploy/local/tests/test_cluster.py
+++ b/mars/deploy/local/tests/test_cluster.py
@@ -155,6 +155,13 @@ def testMultipleOutputTensorExecute(self):
                 np.testing.assert_allclose(U_result, U_expected)
                 np.testing.assert_allclose(s_result, s_expectd)
 
+            with new_session(cluster.endpoint) as session2:
+                U_result, s_result = session2.run(U + 1, s + 1)
+                U_expected, s_expectd, _ = np.linalg.svd(raw, full_matrices=False)
+
+                np.testing.assert_allclose(U_result, U_expected + 1)
+                np.testing.assert_allclose(s_result, s_expectd + 1)
+
     def testGraphFail(self):
         op = SerializeMustFailOperand(f=3)
         tensor = op.new_tensor(None, (3, 3))
diff --git a/mars/tensor/expressions/tests/test_linalg.py b/mars/tensor/expressions/tests/test_linalg.py
--- a/mars/tensor/expressions/tests/test_linalg.py
+++ b/mars/tensor/expressions/tests/test_linalg.py
@@ -77,6 +77,16 @@ def testSVD(self):
         self.assertEqual(len(new_outputs), 3)
         self.assertEqual(len(set([o.op for o in new_outputs])), 1)
 
+        # test tensor graph, do some caculation
+        graph = DirectedGraph()
+        (U + 1).build_graph(tiled=False, graph=graph)
+        (s + 1).build_graph(tiled=False, graph=graph)
+        new_graph = DirectedGraph.from_json(graph.to_json())
+        self.assertEqual((len(new_graph)), 6)
+        new_outputs = [n for n in new_graph if new_graph.count_predecessors(n) == 1]
+        self.assertEqual(len(new_outputs), 5)
+        self.assertEqual(len(set([o.op for o in new_outputs])), 3)
+
     def testLU(self):
         a = mt.random.randint(1, 10, (6, 6), chunks=3)
         p, l, u = mt.linalg.lu(a)