Repository: mars-project/mars

Issue:
[BUG] Failed when remote function's argument is an executed tileable
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**

Failed when remote function's argument is an executed tileable.

**To Reproduce**
To help us reproducing this bug, please provide information below:
1. Your Python version
2. The version of Mars you use
3. Versions of crucial packages, such as numpy, scipy and protobuf
4. Full stack of the error.
5. Minimized code to reproduce the error.

```
In [1]: import mars.dataframe as md                                             

In [2]: from mars.deploy.local import new_cluster                               

In [3]: cluster = new_cluster()                                                 
WARNING: Logging before InitGoogleLogging() is written to STDERR
I0710 12:01:39.413233 286952896 store.cc:1149] Allowing the Plasma store to use up to 3.43597GB of memory.
I0710 12:01:39.414255 286952896 store.cc:1176] Starting object store with directory /tmp and huge page support disabled

In [4]: import mars.tensor as mt                                                

In [5]: df = md.DataFrame(mt.random.rand(10, 3))                                

In [6]: df.execute()                                                            
Out[6]: 
          0         1         2
0  0.212577  0.758511  0.148990
1  0.525289  0.382298  0.331657
2  0.821829  0.991404  0.504818
3  0.910740  0.971152  0.915968
4  0.540863  0.289341  0.546004
5  0.869099  0.257637  0.282307
6  0.738262  0.636345  0.717278
7  0.064604  0.481792  0.356584
8  0.598765  0.156633  0.140831
9  0.873232  0.527147  0.247436

In [7]: def f(in_df): 
   ...:     return in_df.sum().to_pandas() 
   ...:                                                                         

In [8]: import mars.remote as mr                                                

In [9]: mr.spawn(f, args=(df,)).execute()                                       
Unexpected exception occurred in ExecutionActor.execute_graph. graph_key='6ca5f502ce0c9fbccecc434fde3fbe75'
Traceback (most recent call last):
  File "/Users/qinxuye/Workspace/mars/mars/utils.py", line 353, in _wrapped
    return func(*args, **kwargs)
  File "/Users/qinxuye/Workspace/mars/mars/worker/execution.py", line 481, in execute_graph
    no_prepare_chunk_keys=io_meta.get('no_prepare_chunk_keys') or set(),
  File "/Users/qinxuye/Workspace/mars/mars/worker/execution.py", line 57, in __init__
    graph = self.graph = deserialize_graph(graph_serialized)
  File "/Users/qinxuye/Workspace/mars/mars/utils.py", line 283, in deserialize_graph
    return graph_cls.from_pb(g)
  File "mars/graph.pyx", line 426, 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 667, in mars.serialize.core.Serializable.deserialize
  File "mars/serialize/core.pyx", line 810, in mars.serialize.core.Provider.deserialize_model
  File "mars/serialize/core.pyx", line 157, in mars.serialize.core.Field.deserialize
  File "mars/serialize/pbserializer.pyx", line 874, in mars.serialize.pbserializer.ProtobufSerializeProvider.deserialize_field
  File "mars/serialize/pbserializer.pyx", line 839, in mars.serialize.pbserializer.ProtobufSerializeProvider._deserial_reference_value
  File "mars/serialize/core.pyx", line 731, in mars.serialize.core.AttributeAsDict.deserialize
  File "mars/serialize/pbserializer.pyx", line 972, in mars.serialize.pbserializer.ProtobufSerializeProvider.deserialize_attribute_as_dict
  File "mars/serialize/core.pyx", line 157, in mars.serialize.core.Field.deserialize
  File "mars/serialize/pbserializer.pyx", line 878, in mars.serialize.pbserializer.ProtobufSerializeProvider.deserialize_field
  File "mars/serialize/pbserializer.pyx", line 293, in mars.serialize.pbserializer.ProtobufSerializeProvider._get_list
  File "mars/serialize/pbserializer.pyx", line 812, in mars.serialize.pbserializer.ProtobufSerializeProvider._get_value
  File "mars/serialize/pbserializer.pyx", line 802, in mars.serialize.pbserializer.ProtobufSerializeProvider._get_untyped_value
  File "mars/serialize/pbserializer.pyx", line 223, in mars.serialize.pbserializer.ProtobufSerializeProvider._get_function
  File "/Users/qinxuye/Workspace/mars/mars/remote/core.py", line 46, in __setstate__
    chunk = ci[0]().new_chunk(None, _key=ci[1], _id=ci[2], kws=[ci[3]])
  File "/Users/qinxuye/Workspace/mars/mars/operands.py", line 347, in new_chunk
    return self.new_chunks(inputs, kws=kws, **kw)[0]
  File "/Users/qinxuye/Workspace/mars/mars/operands.py", line 341, in new_chunks
    return self._new_chunks(inputs, kws=kws, **kwargs)
  File "/Users/qinxuye/Workspace/mars/mars/dataframe/fetch/core.py", line 39, in _new_chunks
    return super()._new_chunks(inputs, kws=kws, **kw)
  File "/Users/qinxuye/Workspace/mars/mars/operands.py", line 315, in _new_chunks
    chunk = self._create_chunk(j, index, **create_chunk_kw)
  File "/Users/qinxuye/Workspace/mars/mars/operands.py", line 283, in _create_chunk
    raise ValueError('output_type should be specified')
ValueError: output_type should be specified
Unhandled exception in promise call
Traceback (most recent call last):
  File "/Users/qinxuye/Workspace/mars/mars/promise.py", line 372, in _wrapped
    return func(*args, **kwargs)
  File "/Users/qinxuye/Workspace/mars/mars/utils.py", line 353, in _wrapped
    return func(*args, **kwargs)
  File "/Users/qinxuye/Workspace/mars/mars/worker/execution.py", line 481, in execute_graph
    no_prepare_chunk_keys=io_meta.get('no_prepare_chunk_keys') or set(),
  File "/Users/qinxuye/Workspace/mars/mars/worker/execution.py", line 57, in __init__
    graph = self.graph = deserialize_graph(graph_serialized)
  File "/Users/qinxuye/Workspace/mars/mars/utils.py", line 283, in deserialize_graph
    return graph_cls.from_pb(g)
  File "mars/graph.pyx", line 426, 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 667, in mars.serialize.core.Serializable.deserialize
  File "mars/serialize/core.pyx", line 810, in mars.serialize.core.Provider.deserialize_model
  File "mars/serialize/core.pyx", line 157, in mars.serialize.core.Field.deserialize
  File "mars/serialize/pbserializer.pyx", line 874, in mars.serialize.pbserializer.ProtobufSerializeProvider.deserialize_field
  File "mars/serialize/pbserializer.pyx", line 839, in mars.serialize.pbserializer.ProtobufSerializeProvider._deserial_reference_value
  File "mars/serialize/core.pyx", line 731, in mars.serialize.core.AttributeAsDict.deserialize
  File "mars/serialize/pbserializer.pyx", line 972, in mars.serialize.pbserializer.ProtobufSerializeProvider.deserialize_attribute_as_dict
  File "mars/serialize/core.pyx", line 157, in mars.serialize.core.Field.deserialize
  File "mars/serialize/pbserializer.pyx", line 878, in mars.serialize.pbserializer.ProtobufSerializeProvider.deserialize_field
  File "mars/serialize/pbserializer.pyx", line 293, in mars.serialize.pbserializer.ProtobufSerializeProvider._get_list
  File "mars/serialize/pbserializer.pyx", line 812, in mars.serialize.pbserializer.ProtobufSerializeProvider._get_value
  File "mars/serialize/pbserializer.pyx", line 802, in mars.serialize.pbserializer.ProtobufSerializeProvider._get_untyped_value
  File "mars/serialize/pbserializer.pyx", line 223, in mars.serialize.pbserializer.ProtobufSerializeProvider._get_function
  File "/Users/qinxuye/Workspace/mars/mars/remote/core.py", line 46, in __setstate__
    chunk = ci[0]().new_chunk(None, _key=ci[1], _id=ci[2], kws=[ci[3]])
  File "/Users/qinxuye/Workspace/mars/mars/operands.py", line 347, in new_chunk
    return self.new_chunks(inputs, kws=kws, **kw)[0]
  File "/Users/qinxuye/Workspace/mars/mars/operands.py", line 341, in new_chunks
    return self._new_chunks(inputs, kws=kws, **kwargs)
  File "/Users/qinxuye/Workspace/mars/mars/dataframe/fetch/core.py", line 39, in _new_chunks
    return super()._new_chunks(inputs, kws=kws, **kw)
  File "/Users/qinxuye/Workspace/mars/mars/operands.py", line 315, in _new_chunks
    chunk = self._create_chunk(j, index, **create_chunk_kw)
  File "/Users/qinxuye/Workspace/mars/mars/operands.py", line 283, in _create_chunk
    raise ValueError('output_type should be specified')
ValueError: output_type should be specified
Attempt 1: Unexpected error ValueError occurred in executing operand 6ca5f502ce0c9fbccecc434fde3fbe75 in 0.0.0.0:40516
Traceback (most recent call last):
  File "/Users/qinxuye/Workspace/mars/mars/promise.py", line 372, in _wrapped
    return func(*args, **kwargs)
  File "/Users/qinxuye/Workspace/mars/mars/utils.py", line 353, in _wrapped
    return func(*args, **kwargs)
  File "/Users/qinxuye/Workspace/mars/mars/worker/execution.py", line 481, in execute_graph
    no_prepare_chunk_keys=io_meta.get('no_prepare_chunk_keys') or set(),
  File "/Users/qinxuye/Workspace/mars/mars/worker/execution.py", line 57, in __init__
    graph = self.graph = deserialize_graph(graph_serialized)
  File "/Users/qinxuye/Workspace/mars/mars/utils.py", line 283, in deserialize_graph
    return graph_cls.from_pb(g)
  File "mars/graph.pyx", line 426, in mars.graph.DirectedGraph.from_pb
    return cls.deserialize(SerializableGraph.from_pb(pb_obj))
  File "mars/serialize/core.pyx", line 683, in mars.serialize.core.Serializable.from_pb
    return cls.deserialize(provider, obj)
  File "mars/serialize/core.pyx", line 667, in mars.serialize.core.Serializable.deserialize
    obj = provider.deserialize_model(cls, obj, callbacks, key_to_instance)
  File "mars/serialize/core.pyx", line 810, in mars.serialize.core.Provider.deserialize_model
    field.deserialize(self, model_instance, obj, callbacks, key_to_instance)
  File "mars/serialize/core.pyx", line 157, in mars.serialize.core.Field.deserialize
    return provider.deserialize_field(self, model_instance, obj, callbacks, key_to_instance)
  File "mars/serialize/pbserializer.pyx", line 874, in mars.serialize.pbserializer.ProtobufSerializeProvider.deserialize_field
    value = [self._deserial_reference_value(
  File "mars/serialize/pbserializer.pyx", line 839, in mars.serialize.pbserializer.ProtobufSerializeProvider._deserial_reference_value
    return model.deserialize(self, f_obj, callbacks, key_to_instance)
  File "mars/serialize/core.pyx", line 731, in mars.serialize.core.AttributeAsDict.deserialize
    obj = provider.deserialize_attribute_as_dict(
  File "mars/serialize/pbserializer.pyx", line 972, in mars.serialize.pbserializer.ProtobufSerializeProvider.deserialize_attribute_as_dict
    it_field.deserialize(self, model_instance, AttrWrapper(d_obj),
  File "mars/serialize/core.pyx", line 157, in mars.serialize.core.Field.deserialize
    return provider.deserialize_field(self, model_instance, obj, callbacks, key_to_instance)
  File "mars/serialize/pbserializer.pyx", line 878, in mars.serialize.pbserializer.ProtobufSerializeProvider.deserialize_field
    value = self._get_list(field_obj, field.type, callbacks, field.weak_ref)
  File "mars/serialize/pbserializer.pyx", line 293, in mars.serialize.pbserializer.ProtobufSerializeProvider._get_list
    val = self._get_value(it_obj, tp.type if tp is not None else tp,
  File "mars/serialize/pbserializer.pyx", line 812, in mars.serialize.pbserializer.ProtobufSerializeProvider._get_value
    return self._get_untyped_value(obj, callbacks, weak_ref)
  File "mars/serialize/pbserializer.pyx", line 802, in mars.serialize.pbserializer.ProtobufSerializeProvider._get_untyped_value
    return ref(self._get_function(obj))
  File "mars/serialize/pbserializer.pyx", line 223, in mars.serialize.pbserializer.ProtobufSerializeProvider._get_function
    return cloudpickle.loads(x) if x is not None and len(x) > 0 else None
  File "/Users/qinxuye/Workspace/mars/mars/remote/core.py", line 46, in __setstate__
    chunk = ci[0]().new_chunk(None, _key=ci[1], _id=ci[2], kws=[ci[3]])
  File "/Users/qinxuye/Workspace/mars/mars/operands.py", line 347, in new_chunk
    return self.new_chunks(inputs, kws=kws, **kw)[0]
  File "/Users/qinxuye/Workspace/mars/mars/operands.py", line 341, in new_chunks
    return self._new_chunks(inputs, kws=kws, **kwargs)
  File "/Users/qinxuye/Workspace/mars/mars/dataframe/fetch/core.py", line 39, in _new_chunks
    return super()._new_chunks(inputs, kws=kws, **kw)
  File "/Users/qinxuye/Workspace/mars/mars/operands.py", line 315, in _new_chunks
    chunk = self._create_chunk(j, index, **create_chunk_kw)
  File "/Users/qinxuye/Workspace/mars/mars/operands.py", line 283, in _create_chunk
    raise ValueError('output_type should be specified')
ValueError: output_type should be specified
```



Discussion:

Resolving patch:
diff --git a/mars/learn/__init__.py b/mars/learn/__init__.py
--- a/mars/learn/__init__.py
+++ b/mars/learn/__init__.py
@@ -16,12 +16,12 @@
 # import torch first, or some issue emerges,
 # see https://github.com/pytorch/pytorch/issues/2575
 from .contrib import pytorch, tensorflow, xgboost, lightgbm
-from .utils.shuffle import shuffle
 from .metrics import pairwise
 from . import preprocessing
 from . import neighbors
+from . import utils
 
 for _mod in [xgboost, tensorflow, pytorch, lightgbm, neighbors]:
     _mod.register_op()
 
-del _mod, shuffle, pairwise, preprocessing
+del _mod, pairwise, preprocessing, utils
diff --git a/mars/learn/utils/shuffle.py b/mars/learn/utils/shuffle.py
--- a/mars/learn/utils/shuffle.py
+++ b/mars/learn/utils/shuffle.py
@@ -150,8 +150,12 @@ def _calc_chunk_params(cls, in_chunk, axes, chunk_shape, output, output_type,
                 if not no_shuffle and chunk_shape[0] > 1:
                     shape_c[0] = np.nan
             params['shape'] = tuple(shape_c)
-            params['dtypes'] = output.dtypes
-            params['columns_value'] = output.columns_value
+            if 1 not in axes:
+                params['dtypes'] = in_chunk.dtypes
+                params['columns_value'] = in_chunk.columns_value
+            else:
+                params['dtypes'] = output.dtypes
+                params['columns_value'] = output.columns_value
             params['index_value'] = _shuffle_index_value(chunk_op, in_chunk.index_value)
         else:
             assert output_type == OutputType.series
diff --git a/mars/remote/core.py b/mars/remote/core.py
--- a/mars/remote/core.py
+++ b/mars/remote/core.py
@@ -34,19 +34,23 @@ def __init__(self, tileable):
     def __getstate__(self):
         fetch_op = self.tileable.op
         fetch_tileable = self.tileable
-        chunk_infos = [(type(c.op), c.key, c.id, c.params) for c in fetch_tileable.chunks]
-        return type(fetch_op), fetch_op.id, fetch_tileable.params, \
-            fetch_tileable.nsplits, chunk_infos
+        chunk_infos = [(type(c.op), c.op.output_types, c.key, c.id, c.params)
+                       for c in fetch_tileable.chunks]
+        return type(fetch_op), fetch_op.id, fetch_op.output_types, \
+               fetch_tileable.params, fetch_tileable.nsplits, chunk_infos
 
     def __setstate__(self, state):
-        fetch_op_type, fetch_op_id, params, nsplits, chunk_infos = state
+        fetch_op_type, fetch_op_id, output_types, params, nsplits, chunk_infos = state
         params['nsplits'] = nsplits
         chunks = []
         for ci in chunk_infos:
-            chunk = ci[0]().new_chunk(None, _key=ci[1], _id=ci[2], kws=[ci[3]])
+            chunk_op_type, chunk_op_output_types, chunk_key, chunk_id, chunk_params = ci
+            chunk = chunk_op_type(output_types=chunk_op_output_types).new_chunk(
+                None, _key=chunk_key, _id=chunk_id, kws=[chunk_params])
             chunks.append(chunk)
         params['chunks'] = chunks
-        self.tileable = fetch_op_type(_id=fetch_op_id).new_tileable(None, kws=[params])
+        self.tileable = fetch_op_type(
+            _id=fetch_op_id, output_types=output_types).new_tileable(None, kws=[params])
 
     def __mars_tokenize__(self):
         return self.__getstate__()
@@ -219,7 +223,7 @@ def execute(cls, ctx, op: "RemoteFunction"):
                 ctx[out.key] = r
 
 
-def spawn(func, args=(), kwargs=None, retry_when_fail=True, n_output=None):
+def spawn(func, args=(), kwargs=None, retry_when_fail=False, n_output=None):
     """
     Spawn a function and return a Mars Object which can be executed later.
 
@@ -231,7 +235,7 @@ def spawn(func, args=(), kwargs=None, retry_when_fail=True, n_output=None):
        Args to pass to function
     kwargs: dict
        Kwargs to pass to function
-    retry_when_fail: bool, default True
+    retry_when_fail: bool, default False
        If True, retry when function failed.
     n_output: int
        Count of outputs for the function


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
@@ -1076,6 +1076,18 @@ def f3():
             r = session3.run(d, timeout=_exec_timeout)
             np.testing.assert_array_equal(r, np.ones((10, 10)) + 1)
 
+            # test tileable that executed
+            session4 = new_session(cluster.endpoint)
+            df1 = md.DataFrame(raw, chunk_size=3)
+            session4.run(df1)
+
+            def f4(input_df):
+                return input_df.sum().to_pandas()
+
+            d = mr.spawn(f4, args=(df1,), retry_when_fail=False)
+            r = session4.run(d, timeout=_exec_timeout)
+            pd.testing.assert_series_equal(r, pd.DataFrame(raw).sum())
+
     def testLearnInLocalCluster(self, *_):
         from mars.learn.neighbors import NearestNeighbors
         from sklearn.neighbors import NearestNeighbors as SkNearestNeighbors
diff --git a/mars/remote/tests/test_remote_function.py b/mars/remote/tests/test_remote_function.py
--- a/mars/remote/tests/test_remote_function.py
+++ b/mars/remote/tests/test_remote_function.py
@@ -15,8 +15,11 @@
 from collections import defaultdict
 
 import numpy as np
+import pandas as pd
 
+from mars import dataframe as md
 from mars import tensor as mt
+from mars.learn.utils import shuffle
 from mars.remote import spawn, ExecutableTuple
 from mars.session import new_session, Session
 from mars.lib.mmh3 import hash as mmh3_hash
@@ -132,3 +135,18 @@ def f(t, x):
         result = s.execute(session=sess).fetch(session=sess)
         expected = (raw.sum(axis=0) * 3).sum()
         self.assertAlmostEqual(result, expected)
+
+        df1 = md.DataFrame(raw, chunk_size=3)
+        df1.execute(session=sess)
+        df2 = shuffle(df1)
+        df2.execute(session=sess)
+
+        def f2(input_df):
+            return input_df.sum().to_pandas()
+
+        for df in [df1, df2]:
+            s = spawn(f2, args=(df,))
+
+            result = s.execute(session=sess).fetch(session=sess)
+            expected = pd.DataFrame(raw).sum()
+            pd.testing.assert_series_equal(result, expected)