Repository: mars-project/mars

Issue:
[BUG] Failed to run train_test_split in distributed environment
**Describe the bug**
When running ``train_test_split`` in distributed environment, ``'NoneType' object is not callable`` is reported.

**To Reproduce**

```python
rs = np.random.RandomState(0)
df_raw = pd.DataFrame(rs.rand(10, 4))
df = md.DataFrame(df_raw, chunk_size=5)
X, y = df.iloc[:, :-1], df.iloc[:, -1]

for x_to_tensor, y_to_tensor in itertools.product(range(1), range(1)):
    x = X
    if x_to_tensor:
        x = mt.tensor(x)
    yy = y
    if y_to_tensor:
        yy = mt.tensor(yy)

    x_train, x_test, y_train, y_test = train_test_split(
        x, y, random_state=0, session=sess, run_kwargs=run_kwargs)
```

Error log
```
SCH 2021-01-26 14:35:13,477 Mars Scheduler started in standalone mode.
SCH 2021-01-26 14:35:13,477 Actor s:h1:SchedulerClusterInfoActor running in process 72974
SCH 2021-01-26 14:35:13,478 Actor s:h1:ChunkMetaActor running in process 72974 at 127.0.0.1:27341
SCH 2021-01-26 14:35:13,478 Actor s:h1:SessionManagerActor running in process 72974
SCH 2021-01-26 14:35:13,478 Actor s:h1:ResourceActor running in process 72974
SCH 2021-01-26 14:35:13,479 Actor s:h1:NodeInfoActor running in process 72974
WOR1 2021-01-26 14:35:13,484 Setting soft limit to 12.80G.
WOR0 2021-01-26 14:35:13,484 Setting soft limit to 12.80G.
../src/plasma/store.cc:../src/plasma/store.cc1274: Allowing the Plasma store to use up to 0.0104858GB of memory.
:../src/plasma/store.cc:1274: 1297Allowing the Plasma store to use up to : Starting object store with directory /tmp and huge page support disabled
0.0104858GB of memory.
../src/plasma/store.cc:1297: Starting object store with directory /tmp and huge page support disabled
WOR1 2021-01-26 14:35:13,650 Actor w:0:WorkerClusterInfoActor running in process 72983
WOR0 2021-01-26 14:35:13,650 Actor w:0:WorkerClusterInfoActor running in process 72984
WOR0 2021-01-26 14:35:13,654 Actor w:0:WorkerDaemonActor running in process 72984
WOR1 2021-01-26 14:35:13,654 Actor w:0:WorkerDaemonActor running in process 72983
WOR1 2021-01-26 14:35:13,669 Actor w:0:StatusActor running in process 72983
WOR0 2021-01-26 14:35:13,669 Actor w:0:StatusActor running in process 72984
WOR0 2021-01-26 14:35:13,671 Actor w:0:StatusReporterActor running in process 72984
WOR1 2021-01-26 14:35:13,672 Actor w:0:StatusReporterActor running in process 72983
WOR0 2021-01-26 14:35:13,674 Actor w:0:MemQuotaActor running in process 72984
WOR1 2021-01-26 14:35:13,674 Actor w:0:MemQuotaActor running in process 72983
WOR0 2021-01-26 14:35:13,676 Actor w:0:StorageManagerActor running in process 72984
WOR1 2021-01-26 14:35:13,676 Actor w:0:StorageManagerActor running in process 72983
WOR0 2021-01-26 14:35:13,678 Actor w:0:SharedHolderActor running in process 72984
WOR1 2021-01-26 14:35:13,678 Actor w:0:SharedHolderActor running in process 72983
WOR1 2021-01-26 14:35:13,696 Detected actual plasma store size: 10.00M
WOR1 2021-01-26 14:35:13,697 Actor w:0:DispatchActor running in process 72983
WOR0 2021-01-26 14:35:13,697 Detected actual plasma store size: 10.00M
WOR0 2021-01-26 14:35:13,699 Actor w:0:DispatchActor running in process 72984
WOR1 2021-01-26 14:35:13,699 Actor w:0:EventsActor running in process 72983
WOR1 2021-01-26 14:35:13,700 Actor w:0:ReceiverManagerActor running in process 72983
WOR0 2021-01-26 14:35:13,701 Actor w:0:EventsActor running in process 72984
WOR0 2021-01-26 14:35:13,703 Actor w:0:ReceiverManagerActor running in process 72984
WOR1 2021-01-26 14:35:13,703 Actor w:0:ExecutionActor running in process 72983
WOR0 2021-01-26 14:35:13,706 Actor w:0:ExecutionActor running in process 72984
WOR1 2021-01-26 14:35:13,711 Actor w:1:mars-cpu-calc running in process 72986
WOR0 2021-01-26 14:35:13,713 Actor w:1:mars-cpu-calc running in process 72985
WOR1 2021-01-26 14:35:13,722 Slot w:1:mars-cpu-calc registered for queue cpu on 127.0.0.1:58971
WOR0 2021-01-26 14:35:13,723 Slot w:1:mars-cpu-calc registered for queue cpu on 127.0.0.1:49412
WOR1 2021-01-26 14:35:13,726 Actor w:1:mars-inproc-holder running in process 72986
WOR0 2021-01-26 14:35:13,726 Actor w:1:mars-inproc-holder running in process 72985
WOR1 2021-01-26 14:35:13,732 Actor w:1:io_runner_inproc running in process 72986
WOR0 2021-01-26 14:35:13,732 Actor w:1:io_runner_inproc running in process 72985
WOR1 2021-01-26 14:35:13,735 Actor w:2:mars-cpu-calc running in process 72988
WOR0 2021-01-26 14:35:13,738 Actor w:2:mars-cpu-calc running in process 72987
WOR1 2021-01-26 14:35:13,752 Slot w:2:mars-cpu-calc registered for queue cpu on 127.0.0.1:58971
WOR0 2021-01-26 14:35:13,754 Slot w:2:mars-cpu-calc registered for queue cpu on 127.0.0.1:49412
WOR1 2021-01-26 14:35:13,760 Actor w:2:mars-inproc-holder running in process 72988
WOR0 2021-01-26 14:35:13,761 Actor w:2:mars-inproc-holder running in process 72987
WOR1 2021-01-26 14:35:13,767 Actor w:2:io_runner_inproc running in process 72988
WOR0 2021-01-26 14:35:13,768 Actor w:2:io_runner_inproc running in process 72987
WOR1 2021-01-26 14:35:13,773 Actor w:3:mars-sender-0 running in process 72989
WOR0 2021-01-26 14:35:13,776 Actor w:3:mars-sender-0 running in process 72990
WOR0 2021-01-26 14:35:13,795 Slot w:3:mars-sender-0 registered for queue sender on 127.0.0.1:49412
WOR1 2021-01-26 14:35:13,795 Slot w:3:mars-sender-0 registered for queue sender on 127.0.0.1:58971
WOR0 2021-01-26 14:35:13,799 Actor w:4:mars-sender-1 running in process 72991
WOR1 2021-01-26 14:35:13,800 Actor w:4:mars-sender-1 running in process 72992
WOR0 2021-01-26 14:35:13,814 Slot w:4:mars-sender-1 registered for queue sender on 127.0.0.1:49412
WOR1 2021-01-26 14:35:13,815 Slot w:4:mars-sender-1 registered for queue sender on 127.0.0.1:58971
WOR0 2021-01-26 14:35:13,818 Actor w:5:mars-sender-2 running in process 72993
WOR1 2021-01-26 14:35:13,819 Actor w:5:mars-sender-2 running in process 72994
WOR0 2021-01-26 14:35:13,834 Slot w:5:mars-sender-2 registered for queue sender on 127.0.0.1:49412
WOR1 2021-01-26 14:35:13,835 Slot w:5:mars-sender-2 registered for queue sender on 127.0.0.1:58971
WOR0 2021-01-26 14:35:13,838 Actor w:6:mars-sender-3 running in process 72996
WOR1 2021-01-26 14:35:13,841 Actor w:6:mars-sender-3 running in process 72995
WOR0 2021-01-26 14:35:13,854 Slot w:6:mars-sender-3 registered for queue sender on 127.0.0.1:49412
WOR1 2021-01-26 14:35:13,855 Slot w:6:mars-sender-3 registered for queue sender on 127.0.0.1:58971
WOR0 2021-01-26 14:35:13,857 Slot w:3:mars-custom-log-fetch registered for queue custom_log on 127.0.0.1:49412
WOR1 2021-01-26 14:35:13,858 Slot w:3:mars-custom-log-fetch registered for queue custom_log on 127.0.0.1:58971
WOR0 2021-01-26 14:35:13,861 Slot w:4:mars-custom-log-fetch registered for queue custom_log on 127.0.0.1:49412
WOR1 2021-01-26 14:35:13,863 Slot w:4:mars-custom-log-fetch registered for queue custom_log on 127.0.0.1:58971
WOR0 2021-01-26 14:35:13,864 Slot w:5:mars-custom-log-fetch registered for queue custom_log on 127.0.0.1:49412
WOR1 2021-01-26 14:35:13,865 Slot w:5:mars-custom-log-fetch registered for queue custom_log on 127.0.0.1:58971
WOR0 2021-01-26 14:35:13,867 Slot w:6:mars-custom-log-fetch registered for queue custom_log on 127.0.0.1:49412
WOR1 2021-01-26 14:35:13,868 Slot w:6:mars-custom-log-fetch registered for queue custom_log on 127.0.0.1:58971
WOR0 2021-01-26 14:35:13,869 Actor w:3:mars-receiver-0 running in process 72990
WOR1 2021-01-26 14:35:13,870 Actor w:3:mars-receiver-0 running in process 72989
WOR0 2021-01-26 14:35:13,873 Slot w:3:mars-receiver-0 registered for queue receiver on 127.0.0.1:49412
WOR1 2021-01-26 14:35:13,873 Slot w:3:mars-receiver-0 registered for queue receiver on 127.0.0.1:58971
WOR0 2021-01-26 14:35:13,875 Actor w:3:mars-receiver-1 running in process 72990
WOR1 2021-01-26 14:35:13,875 Actor w:3:mars-receiver-1 running in process 72989
WOR0 2021-01-26 14:35:13,878 Slot w:3:mars-receiver-1 registered for queue receiver on 127.0.0.1:49412
WOR1 2021-01-26 14:35:13,878 Slot w:3:mars-receiver-1 registered for queue receiver on 127.0.0.1:58971
WOR0 2021-01-26 14:35:13,879 Actor w:4:mars-receiver-2 running in process 72991
WOR1 2021-01-26 14:35:13,880 Actor w:4:mars-receiver-2 running in process 72992
WOR0 2021-01-26 14:35:13,883 Slot w:4:mars-receiver-2 registered for queue receiver on 127.0.0.1:49412
WOR1 2021-01-26 14:35:13,884 Slot w:4:mars-receiver-2 registered for queue receiver on 127.0.0.1:58971
WOR0 2021-01-26 14:35:13,885 Actor w:4:mars-receiver-3 running in process 72991
WOR1 2021-01-26 14:35:13,885 Actor w:4:mars-receiver-3 running in process 72992
WOR0 2021-01-26 14:35:13,887 Slot w:4:mars-receiver-3 registered for queue receiver on 127.0.0.1:49412
WOR1 2021-01-26 14:35:13,888 Slot w:4:mars-receiver-3 registered for queue receiver on 127.0.0.1:58971
WOR0 2021-01-26 14:35:13,889 Actor w:5:mars-receiver-4 running in process 72993
WOR1 2021-01-26 14:35:13,890 Actor w:5:mars-receiver-4 running in process 72994
WOR0 2021-01-26 14:35:13,892 Slot w:5:mars-receiver-4 registered for queue receiver on 127.0.0.1:49412
WOR1 2021-01-26 14:35:13,893 Slot w:5:mars-receiver-4 registered for queue receiver on 127.0.0.1:58971
WOR0 2021-01-26 14:35:13,894 Actor w:5:mars-receiver-5 running in process 72993
WOR1 2021-01-26 14:35:13,895 Actor w:5:mars-receiver-5 running in process 72994
WOR0 2021-01-26 14:35:13,897 Slot w:5:mars-receiver-5 registered for queue receiver on 127.0.0.1:49412
WOR1 2021-01-26 14:35:13,898 Slot w:5:mars-receiver-5 registered for queue receiver on 127.0.0.1:58971
WOR0 2021-01-26 14:35:13,899 Actor w:6:mars-receiver-6 running in process 72996
WOR1 2021-01-26 14:35:13,899 Actor w:6:mars-receiver-6 running in process 72995
WOR0 2021-01-26 14:35:13,902 Slot w:6:mars-receiver-6 registered for queue receiver on 127.0.0.1:49412
WOR1 2021-01-26 14:35:13,903 Slot w:6:mars-receiver-6 registered for queue receiver on 127.0.0.1:58971
WOR0 2021-01-26 14:35:13,904 Actor w:6:mars-receiver-7 running in process 72996
WOR1 2021-01-26 14:35:13,904 Actor w:6:mars-receiver-7 running in process 72995
WOR0 2021-01-26 14:35:13,908 Slot w:6:mars-receiver-7 registered for queue receiver on 127.0.0.1:49412
WOR1 2021-01-26 14:35:13,908 Slot w:6:mars-receiver-7 registered for queue receiver on 127.0.0.1:58971
WOR1 2021-01-26 14:35:13,910 Actor w:0:mars-process-helper running in process 72983
WOR0 2021-01-26 14:35:13,910 Actor w:0:mars-process-helper running in process 72984
WOR1 2021-01-26 14:35:13,911 Slot w:0:mars-process-helper registered for queue process_helper on 127.0.0.1:58971
WOR0 2021-01-26 14:35:13,911 Slot w:0:mars-process-helper registered for queue process_helper on 127.0.0.1:49412
WOR1 2021-01-26 14:35:13,912 Actor w:1:mars-process-helper running in process 72986
WOR0 2021-01-26 14:35:13,913 Actor w:1:mars-process-helper running in process 72985
WOR0 2021-01-26 14:35:13,915 Slot w:1:mars-process-helper registered for queue process_helper on 127.0.0.1:49412
WOR1 2021-01-26 14:35:13,915 Slot w:1:mars-process-helper registered for queue process_helper on 127.0.0.1:58971
WOR1 2021-01-26 14:35:13,918 Actor w:2:mars-process-helper running in process 72988
WOR0 2021-01-26 14:35:13,918 Actor w:2:mars-process-helper running in process 72987
WOR0 2021-01-26 14:35:13,921 Slot w:2:mars-process-helper registered for queue process_helper on 127.0.0.1:49412
WOR1 2021-01-26 14:35:13,921 Slot w:2:mars-process-helper registered for queue process_helper on 127.0.0.1:58971
WOR0 2021-01-26 14:35:13,924 Actor w:3:mars-process-helper running in process 72990
WOR1 2021-01-26 14:35:13,924 Actor w:3:mars-process-helper running in process 72989
WOR0 2021-01-26 14:35:13,926 Slot w:3:mars-process-helper registered for queue process_helper on 127.0.0.1:49412
WOR1 2021-01-26 14:35:13,926 Slot w:3:mars-process-helper registered for queue process_helper on 127.0.0.1:58971
WOR0 2021-01-26 14:35:13,930 Actor w:4:mars-process-helper running in process 72991
WOR1 2021-01-26 14:35:13,930 Actor w:4:mars-process-helper running in process 72992
WOR0 2021-01-26 14:35:13,937 Slot w:4:mars-process-helper registered for queue process_helper on 127.0.0.1:49412
WOR1 2021-01-26 14:35:13,939 Slot w:4:mars-process-helper registered for queue process_helper on 127.0.0.1:58971
WOR0 2021-01-26 14:35:13,944 Actor w:5:mars-process-helper running in process 72993
WOR1 2021-01-26 14:35:13,945 Actor w:5:mars-process-helper running in process 72994
WOR0 2021-01-26 14:35:13,947 Slot w:5:mars-process-helper registered for queue process_helper on 127.0.0.1:49412
WOR1 2021-01-26 14:35:13,949 Slot w:5:mars-process-helper registered for queue process_helper on 127.0.0.1:58971
WOR0 2021-01-26 14:35:13,954 Actor w:6:mars-process-helper running in process 72996
WOR1 2021-01-26 14:35:13,954 Actor w:6:mars-process-helper running in process 72995
WOR0 2021-01-26 14:35:13,958 Slot w:6:mars-process-helper registered for queue process_helper on 127.0.0.1:49412
WOR1 2021-01-26 14:35:13,959 Slot w:6:mars-process-helper registered for queue process_helper on 127.0.0.1:58971
WOR0 2021-01-26 14:35:13,961 Actor w:0:ResultSenderActor running in process 72984
WOR1 2021-01-26 14:35:13,962 Actor w:0:ResultSenderActor running in process 72983
WOR0 2021-01-26 14:35:13,963 Actor w:0:ResultCopyActor running in process 72984
WOR1 2021-01-26 14:35:13,964 Actor w:0:ResultCopyActor running in process 72983
WARNING:bokeh.server.util:Host wildcard '*' will allow connections originating from multiple (or possibly all) hostnames or IPs. Use non-wildcard values to restrict access explicitly
WEB 2021-01-26 14:35:20,232 Mars UI started at 127.0.0.1:21288
SCH 2021-01-26 14:35:20,874 Actor s:h1:session$b5582d85e3fdfe2b0538c3f4a6751fb5 running in process 72974
SCH 2021-01-26 14:35:20,875 Actor s:h1:assigner$b5582d85e3fdfe2b0538c3f4a6751fb5 running in process 72974
SCH 2021-01-26 14:35:20,875 Actor s:0:AssignEvaluationActor$b5582d85e3fdfe2b0538c3f4a6751fb5 running in process 72974
SCH 2021-01-26 14:35:20,929 Actor s:0:graph$b5582d85e3fdfe2b0538c3f4a6751fb5$9bb221d1340d2ccc125c01ed42330c23 running in process 72974
SCH 2021-01-26 14:35:20,930 Graph 9bb221d1340d2ccc125c01ed42330c23 state from GraphState.UNSCHEDULED to GraphState.PREPARING.
SCH 2021-01-26 14:35:20,946 Begin preparing graph 9bb221d1340d2ccc125c01ed42330c23 with 9 tileables to chunk graph.
SCH 2021-01-26 14:35:20,968 Terminal chunk keys: {'ef47ef4796e60445c544df4932449c98', 'e5d3a577f04edd5ed13e4d1bbedd6835', '5aef6cf6563ec183673159b483273a69', '965705674a7b15205fbe9f303957fe2a', '113cf9a2d1747ee1d168588b6416755f', 'e319d64c44060fb66e67d1c6770d43f9'}
SCH 2021-01-26 14:35:20,971 Placing initial chunks for graph 9bb221d1340d2ccc125c01ed42330c23
SCH 2021-01-26 14:35:20,972 Worker assign quotas: {'127.0.0.1:49412': 1, '127.0.0.1:58971': 1}
SCH 2021-01-26 14:35:20,972 Scan spread ranges: {}
SCH 2021-01-26 14:35:20,973 Creating operand actors for graph 9bb221d1340d2ccc125c01ed42330c23
SCH 2021-01-26 14:35:20,979 Operand actor creation progress: 1 / 18
SCH 2021-01-26 14:35:20,982 Operand actor creation progress: 2 / 18
SCH 2021-01-26 14:35:20,983 Unexpected exception occurred in GraphActor.get_executable_operand_dag. op_key=241a38623560db72c868622ae4cf87da
Traceback (most recent call last):
  File "~/Code/mars/mars/scheduler/graph.py", line 830, in get_executable_operand_dag
    inp_chunk = input_mapping[(inp.key, inp.id)]
KeyError: ('8501fa868db396f18599c9990f1f1c2a', '140271835165984')

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "~/Code/mars/mars/utils.py", line 383, in _wrapped
    return func(*args, **kwargs)
  File "~/Code/mars/mars/scheduler/graph.py", line 833, in get_executable_operand_dag
    = build_fetch_chunk(inp, input_chunk_keys).data
  File "~/Code/mars/mars/utils.py", line 581, in build_fetch_chunk
    op = chunk_op.get_fetch_op_cls(chunk)(to_fetch_keys=to_fetch_keys, to_fetch_idxes=to_fetch_idxes)
  File "~/Code/mars/mars/operands.py", line 551, in _inner
    return cls(output_types=output_types, **kw)
TypeError: 'NoneType' object is not callable
Traceback (most recent call last):
  File "~/Code/mars/mars/scheduler/graph.py", line 830, in get_executable_operand_dag
    inp_chunk = input_mapping[(inp.key, inp.id)]
KeyError: ('8501fa868db396f18599c9990f1f1c2a', '140271835165984')

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "~/miniconda3/lib/python3.8/site-packages/gevent/threadpool.py", line 167, in __run_task
    thread_result.set(func(*args, **kwargs))
  File "mars/actors/pool/gevent_pool.pyx", line 171, in mars.actors.pool.gevent_pool.GeventThreadPool._wrap_watch_sem.inner
  File "~/Code/mars/mars/utils.py", line 383, in _wrapped
    return func(*args, **kwargs)
  File "~/Code/mars/mars/scheduler/graph.py", line 833, in get_executable_operand_dag
    = build_fetch_chunk(inp, input_chunk_keys).data
  File "~/Code/mars/mars/utils.py", line 581, in build_fetch_chunk
    op = chunk_op.get_fetch_op_cls(chunk)(to_fetch_keys=to_fetch_keys, to_fetch_idxes=to_fetch_idxes)
  File "~/Code/mars/mars/operands.py", line 551, in _inner
    return cls(output_types=output_types, **kw)
TypeError: 'NoneType' object is not callable
2021-01-26T06:35:20Z (<ThreadPoolWorker at 0x7f939500b930 thread_ident=0x700007417000 threadpool-hub=<Hub at 0x7f9394dbad00 thread_ident=0x119739dc0>>, <cyfunction GeventThreadPool._wrap_watch_sem.<locals>.inner at 0x7f9394ea2790>) failed with TypeError
```


Discussion:

Resolving patch:
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
@@ -253,7 +253,8 @@ def tile(cls, op):
                     map_chunk = map_chunk_op.new_chunk([in_chunk], **params)
                     map_chunks.append(map_chunk)
 
-                proxy_chunk = LearnShuffleProxy(_tensor_keys=[inp.key]).new_chunk(map_chunks)
+                proxy_chunk = LearnShuffleProxy(
+                    _tensor_keys=[inp.key], output_types=[output_type]).new_chunk(map_chunks)
 
                 reduce_axes = tuple(ax for j, ax in enumerate(inp_axes) if reduce_sizes[j] > 1)
                 reduce_sizes_ = tuple(rs for rs in reduce_sizes if rs > 1)


Test patch:
diff --git a/mars/learn/model_selection/tests/integrated/__init__.py b/mars/learn/model_selection/tests/integrated/__init__.py
new file mode 100644
--- /dev/null
+++ b/mars/learn/model_selection/tests/integrated/__init__.py
@@ -0,0 +1,13 @@
+# Copyright 1999-2020 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
diff --git a/mars/learn/model_selection/tests/integrated/test_distributed_split.py b/mars/learn/model_selection/tests/integrated/test_distributed_split.py
new file mode 100644
--- /dev/null
+++ b/mars/learn/model_selection/tests/integrated/test_distributed_split.py
@@ -0,0 +1,62 @@
+# Copyright 1999-2020 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import itertools
+import os
+import sys
+import unittest
+
+import numpy as np
+import pandas as pd
+
+try:
+    import sklearn
+except ImportError:
+    sklearn = None
+
+import mars.dataframe as md
+import mars.tensor as mt
+from mars.learn.model_selection import train_test_split
+from mars.session import new_session
+from mars.tests.integrated.base import IntegrationTestBase
+
+
+@unittest.skipIf(sklearn is None, 'sklearn not installed')
+@unittest.skipIf(sys.platform == 'win32', "plasma don't support windows")
+class Test(IntegrationTestBase):
+    def testDistributedSplit(self):
+        service_ep = 'http://127.0.0.1:' + self.web_port
+        timeout = 120 if 'CI' in os.environ else -1
+        with new_session(service_ep) as sess:
+            run_kwargs = {'timeout': timeout}
+
+            rs = np.random.RandomState(0)
+            df_raw = pd.DataFrame(rs.rand(10, 4))
+            df = md.DataFrame(df_raw, chunk_size=5)
+            X, y = df.iloc[:, :-1], df.iloc[:, -1]
+
+            for x_to_tensor, y_to_tensor in itertools.product(range(1), range(1)):
+                x = X
+                if x_to_tensor:
+                    x = mt.tensor(x)
+                yy = y
+                if y_to_tensor:
+                    yy = mt.tensor(yy)
+
+                x_train, x_test, y_train, y_test = train_test_split(
+                    x, y, random_state=0, session=sess, run_kwargs=run_kwargs)
+                self.assertIsInstance(x_train, type(x))
+                self.assertIsInstance(x_test, type(x))
+                self.assertIsInstance(y_train, type(yy))
+                self.assertIsInstance(y_test, type(yy))