Repository: mars-project/mars

Issue:
[BUG] Run KNN failed in cluster mode
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**
Create a local cluster and run NearestNeighbors failed.

**To Reproduce**
``` Python
import numpy as np
import mars.tensor as mt
from mars.learn.neighbors import NearestNeighbors
from mars.deploy.local import new_cluster
with new_cluster(scheduler_n_process=2, worker_n_process=2, shared_memory='20M', web=False) as cluster:
    rs = np.random.RandomState(0)
    raw_X = rs.rand(10, 5)
    raw_Y = rs.rand(8, 5)

    X = mt.tensor(raw_X, chunk_size=7)
    Y = mt.tensor(raw_Y, chunk_size=(5, 3))
    nn = NearestNeighbors(n_neighbors=3)
    nn.fit(X)

-----------------------------------------------------------------------------------------
Traceback (most recent call last):
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/utils.py", line 339, in _wrapped
    return func(*args, **kwargs)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/scheduler/graph.py", line 436, in get_chunk_graph
    return self._chunk_graph_builder.iterative_chunk_graphs[-1]
IndexError: list index out of range
Unexpected exception occurred in GraphActor.stop_graph.
Traceback (most recent call last):
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/scheduler/graph.py", line 371, in _execute_graph
    self.prepare_graph(compose=compose)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/utils.py", line 339, in _wrapped
    return func(*args, **kwargs)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/utils.py", line 473, in inner
    return func(*args, **kwargs)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/scheduler/graph.py", line 600, in prepare_graph
    self._target_tileable_datas + fetch_tileables, tileable_graph)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/utils.py", line 385, in _wrapped
    return func(*args, **kwargs)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/utils.py", line 473, in inner
    return func(*args, **kwargs)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/tiles.py", line 342, in build
    tileables, tileable_graph=tileable_graph)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/utils.py", line 385, in _wrapped
    return func(*args, **kwargs)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/utils.py", line 473, in inner
    return func(*args, **kwargs)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/tiles.py", line 255, in build
    self._on_tile_failure(tileable_data.op, exc_info)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/tiles.py", line 294, in inner
    raise exc_info[1].with_traceback(exc_info[2]) from None
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/tiles.py", line 235, in build
    tiled = self._tile(tileable_data, tileable_graph)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/tiles.py", line 330, in _tile
    return super()._tile(tileable_data, tileable_graph)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/tiles.py", line 195, in _tile
    tds = on_tile(tileable_data.op.outputs, tds)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/scheduler/graph.py", line 581, in on_tile
    return [self.tile_fetch_tileable(first)]
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/scheduler/graph.py", line 1139, in tile_fetch_tileable
    fetch_graph = deserialize_graph(graph_ref.build_fetch_graph(tileable_key))
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/utils.py", line 272, in deserialize_graph
    return graph_cls.from_pb(g)
  File "mars/graph.pyx", line 440, in mars.graph.DirectedGraph.from_pb
  File "mars/graph.pyx", line 437, in mars.graph.DirectedGraph.from_pb
  File "mars/serialize/core.pyx", line 683, in mars.serialize.core.Serializable.from_pb
  File "mars/serialize/core.pyx", line 669, in mars.serialize.core.Serializable.deserialize
  File "mars/serialize/pbserializer.pyx", line 876, in mars.serialize.pbserializer.ProtobufSerializeProvider.deserialize_field.cb
KeyError: ('9a6fa72cf124c991c7f75a489b4c1b13', '5299152656')
```



Discussion:

Resolving patch:
diff --git a/mars/scheduler/graph.py b/mars/scheduler/graph.py
--- a/mars/scheduler/graph.py
+++ b/mars/scheduler/graph.py
@@ -1126,7 +1126,7 @@ def build_fetch_graph(self, tileable_key):
         tileable = self._get_tileable_by_key(tileable_key)
         graph = DAG()
 
-        new_tileable = build_fetch_tileable(tileable)
+        new_tileable = build_fetch_tileable(tileable).data
         graph.add_node(new_tileable)
         return serialize_graph(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
@@ -980,3 +980,27 @@ def f2():
 
             r = session2.run(d, timeout=_exec_timeout)
             self.assertEqual(r, 6)
+
+    def testKNNInLocalCluster(self, *_):
+        from mars.learn.neighbors import NearestNeighbors
+        from sklearn.neighbors import NearestNeighbors as SkNearestNeighbors
+
+        with new_cluster(scheduler_n_process=2, worker_n_process=3, shared_memory='20M') as cluster:
+            rs = np.random.RandomState(0)
+            raw_X = rs.rand(10, 5)
+            raw_Y = rs.rand(8, 5)
+
+            X = mt.tensor(raw_X, chunk_size=7)
+            Y = mt.tensor(raw_Y, chunk_size=(5, 3))
+            nn = NearestNeighbors(n_neighbors=3)
+            nn.fit(X)
+
+            ret = nn.kneighbors(Y, session=cluster.session)
+
+            snn = SkNearestNeighbors(n_neighbors=3)
+            snn.fit(raw_X)
+            expected = snn.kneighbors(raw_Y)
+
+            result = [r.fetch() for r in ret]
+            np.testing.assert_almost_equal(result[0], expected[0])
+            np.testing.assert_almost_equal(result[1], expected[1])