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.

Resolving patch:
diff --git a/mars/operands/random.py b/mars/operands/random.py
--- a/mars/operands/random.py
+++ b/mars/operands/random.py
@@ -27,7 +27,8 @@
 
 
 class RandomOperand(Operand):
-    _state = TupleField('state')
+    _state = TupleField('state', on_serialize=lambda x: tuple(x) if x is not None else x,
+                        on_deserialize=lambda x: State(x) if x is not None else x)
 
     @property
     def state(self):
diff --git a/mars/tensor/core.py b/mars/tensor/core.py
--- a/mars/tensor/core.py
+++ b/mars/tensor/core.py
@@ -290,7 +290,7 @@ def build_graph(self, graph=None, cls=DAG, tiled=False, compose=True):
             nodes = list(c.data for c in self.chunks)
             keys = list(c.key for c in self.chunks)
         else:
-            nodes = [self]
+            nodes = list(self.op.outputs)
         visited = set()
         while len(nodes) > 0:
             chunk = nodes.pop()


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
@@ -90,7 +90,7 @@ def testNSchedulersNWorkers(self):
         self.assertEqual(LocalDistributedCluster._calc_scheduler_worker_n_process(
             5, 3, 2, calc_cpu_count=calc_cpu_cnt), (3, 2))
 
-    def testTensorExecute(self):
+    def testSingleOutputTensorExecute(self):
         with new_cluster(scheduler_n_process=2, worker_n_process=2) as cluster:
             self.assertIs(cluster.session, Session.default_or_local())
 
@@ -106,3 +106,36 @@ def testTensorExecute(self):
 
             res = r.execute()
             self.assertLess(res, 39)
+
+    def testMultipleOutputTensorExecute(self):
+        with new_cluster(scheduler_n_process=2, worker_n_process=2) as cluster:
+            session = cluster.session
+
+            t = mt.random.rand(20, 5, chunks=5)
+            r = mt.linalg.svd(t)
+
+            res = session.run((t,) + r)
+
+            U, s, V = res[1:]
+            np.testing.assert_allclose(res[0], U.dot(np.diag(s).dot(V)))
+
+            raw = np.random.rand(20, 5)
+
+            # to test the fuse, the graph should be fused
+            t = mt.array(raw)
+            U, s, V = mt.linalg.svd(t)
+            r = U.dot(mt.diag(s).dot(V))
+
+            res = r.execute()
+            np.testing.assert_allclose(raw, res)
+
+            # test submit part of svd outputs
+            t = mt.array(raw)
+            U, s, V = mt.linalg.svd(t)
+
+            with new_session(cluster.endpoint) as session2:
+                U_result, s_result = session2.run(U, s)
+                U_expected, s_expectd, _ = np.linalg.svd(raw, full_matrices=False)
+
+                np.testing.assert_allclose(U_result, U_expected)
+                np.testing.assert_allclose(s_result, s_expectd)
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
@@ -19,6 +19,7 @@
 import numpy as np
 
 import mars.tensor as mt
+from mars.graph import DirectedGraph
 
 
 class Test(unittest.TestCase):
@@ -62,6 +63,20 @@ def testSVD(self):
         self.assertEqual(len(s.chunks), 1)
         self.assertEqual(len(V.chunks), 1)
 
+        rs = mt.random.RandomState(1)
+        a = rs.rand(9, 6, chunks=(3, 6))
+        U, s, V = mt.linalg.svd(a)
+
+        # test tensor graph
+        graph = DirectedGraph()
+        U.build_graph(tiled=False, graph=graph)
+        s.build_graph(tiled=False, graph=graph)
+        new_graph = DirectedGraph.from_json(graph.to_json())
+        self.assertEqual((len(new_graph)), 4)
+        new_outputs = [n for n in new_graph if new_graph.count_predecessors(n) == 1]
+        self.assertEqual(len(new_outputs), 3)
+        self.assertEqual(len(set([o.op for o in new_outputs])), 1)
+
     def testLU(self):
         a = mt.random.randint(1, 10, (6, 6), chunks=3)
         p, l, u = mt.linalg.lu(a)