We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[BUG] DASK-on-Mars failed when compute multiple objects
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**

DASK-on-Mars failed when compute multiple objects.

**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 pandas
4. Full stack of the error.
5. Minimized code to reproduce the error.

```python
import dask
import mars
from mars.contrib.dask import mars_scheduler

def add(x: int):
    return x + 1


mars.new_session()
xs = [dask.delayed(add)(i) for i in range(10)]
dask.compute(xs, scheduler=mars_scheduler)
```

Output:

```
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-12-531953a29902> in <module>
      4 
      5 xs = [dask.delayed(add)(i) for i in range(10)]
----> 6 dask.compute(xs, scheduler=mars_scheduler)

~/miniconda3/lib/python3.7/site-packages/dask/base.py in compute(*args, **kwargs)
    567 
    568     results = schedule(dsk, keys, **kwargs)
--> 569     return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])
    570 
    571 

~/miniconda3/lib/python3.7/site-packages/dask/base.py in repack(results)
    448         dsk = repack_dsk.copy()
    449         dsk[collections_token] = quote(results)
--> 450         return simple_get(dsk, out)
    451 
    452     return collections, repack

~/miniconda3/lib/python3.7/site-packages/dask/core.py in get(dsk, out, cache)
    149     for key in toposort(dsk):
    150         task = dsk[key]
--> 151         result = _execute_task(task, cache)
    152         cache[key] = result
    153     result = _execute_task(out, cache)

~/miniconda3/lib/python3.7/site-packages/dask/core.py in _execute_task(arg, cache, dsk)
    119         # temporaries by their reference count and can execute certain
    120         # operations in-place.
--> 121         return func(*(_execute_task(a, cache) for a in args))
    122     elif not ishashable(arg):
    123         return arg

IndexError: list index out of range
```

[BUG] DASK-on-Mars failed when compute multiple objects
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**

DASK-on-Mars failed when compute multiple objects.

**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 pandas
4. Full stack of the error.
5. Minimized code to reproduce the error.

```python
import dask
import mars
from mars.contrib.dask import mars_scheduler

def add(x: int):
    return x + 1


mars.new_session()
xs = [dask.delayed(add)(i) for i in range(10)]
dask.compute(xs, scheduler=mars_scheduler)
```

Output:

```
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-12-531953a29902> in <module>
      4 
      5 xs = [dask.delayed(add)(i) for i in range(10)]
----> 6 dask.compute(xs, scheduler=mars_scheduler)

~/miniconda3/lib/python3.7/site-packages/dask/base.py in compute(*args, **kwargs)
    567 
    568     results = schedule(dsk, keys, **kwargs)
--> 569     return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])
    570 
    571 

~/miniconda3/lib/python3.7/site-packages/dask/base.py in repack(results)
    448         dsk = repack_dsk.copy()
    449         dsk[collections_token] = quote(results)
--> 450         return simple_get(dsk, out)
    451 
    452     return collections, repack

~/miniconda3/lib/python3.7/site-packages/dask/core.py in get(dsk, out, cache)
    149     for key in toposort(dsk):
    150         task = dsk[key]
--> 151         result = _execute_task(task, cache)
    152         cache[key] = result
    153     result = _execute_task(out, cache)

~/miniconda3/lib/python3.7/site-packages/dask/core.py in _execute_task(arg, cache, dsk)
    119         # temporaries by their reference count and can execute certain
    120         # operations in-place.
--> 121         return func(*(_execute_task(a, cache) for a in args))
    122     elif not ishashable(arg):
    123         return arg

IndexError: list index out of range
```


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `mars/contrib/dask/utils.py`
Content:
```
1 # Copyright 1999-2021 Alibaba Group Holding Ltd.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #      http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 from typing import List
16 
17 from dask import is_dask_collection
18 from dask.array.core import _concatenate2 as array_concat
19 from dask.dataframe import concat as df_concat
20 from dask.utils import is_arraylike, is_dataframe_like, is_series_like, is_index_like
21 
22 from ...remote import spawn
23 
24 
25 def concat(objs: List):
26     """
27     Concat the results of partitioned dask task executions. This function guess the
28         types of resulting list, then calls the corresponding native dask concat functions.
29 
30     Parameters
31     ----------
32     objs: List
33         List of the partitioned dask task execution results, which will be concat.
34 
35     Returns
36     -------
37     obj:
38         The concat result
39 
40     """
41     if is_arraylike(objs[0]):
42         res = array_concat(objs, axes=[0])  # TODO: Add concat with args support
43     elif any((is_dataframe_like(objs[0]), is_series_like(objs[0]), is_index_like(objs[0]))):
44         res = df_concat(objs)
45     else:
46         res = objs
47         while isinstance(res, List):
48             res = res[0]
49     return res.compute() if is_dask_collection(res) else res
50 
51 
52 def reduce(objs: List[List]):
53     """
54     Spawn a concat task for 2d-list objects
55 
56     Parameters
57     ----------
58     objs: List
59         2d-list of the partitioned dask task execution results, which will be concat.
60 
61     Returns
62     -------
63     obj:
64         The spawning concat task.
65     """
66     return spawn(
67         concat,
68         args=([spawn(concat, args=(objs_d,)) for objs_d in objs],),
69     )
70
```

Path: `mars/contrib/dask/scheduler.py`
Content:
```
1 # Copyright 1999-2021 Alibaba Group Holding Ltd.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #      http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 from typing import List, Tuple
16 
17 from dask.core import istask, ishashable
18 
19 from .utils import reduce
20 from ...remote import spawn
21 
22 
23 def mars_scheduler(dsk: dict, keys: List[List[str]]):
24     """
25     A Dask-Mars scheduler
26 
27     This scheduler is intended to be compatible with existing
28     dask user interface, no callbacks are implemented.
29 
30     Parameters
31     ----------
32     dsk: Dict
33         Dask graph, represented as a task DAG dictionary.
34     keys: List[List[str]]
35         2d-list of Dask graph keys whose values we wish to compute and return.
36 
37     Returns
38     -------
39     Object
40         Computed values corresponding to the provided keys.
41     """
42 
43     return [[reduce(mars_dask_get(dsk, keys)).execute().fetch()]]
44 
45 
46 def mars_dask_get(dsk: dict, keys: List[List]):
47     """
48     A Dask-Mars convert function. This function will send the dask graph layers
49         to Mars Remote API, generating mars objects correspond to the provided keys.
50 
51     Parameters
52     ----------
53     dsk: Dict
54         Dask graph, represented as a task DAG dictionary.
55     keys: List[List[str]]
56         2d-list of Dask graph keys whose values we wish to compute and return.
57 
58     Returns
59     -------
60     Object
61         Spawned mars objects corresponding to the provided keys.
62     """
63 
64     def _get_arg(a):
65         # if arg contains layer index or callable objs, handle it
66         if ishashable(a) and a in dsk.keys():
67             while ishashable(a) and a in dsk.keys():
68                 a = dsk[a]
69             return _execute_task(a)
70         elif not isinstance(a, str) and hasattr(a, "__getitem__"):
71             if istask(a):  # TODO:Handle `SubgraphCallable`, which may contains dsk in it
72                 return spawn(a[0], args=tuple(_get_arg(i) for i in a[1:]))
73             elif isinstance(a, dict):
74                 return {k: _get_arg(v) for k, v in a.items()}
75             elif isinstance(a, List) or isinstance(a, Tuple):
76                 return type(a)(_get_arg(i) for i in a)
77         return a
78 
79     def _execute_task(task: tuple):
80         if not istask(task):
81             return _get_arg(task)
82         return spawn(task[0], args=tuple(_get_arg(a) for a in task[1:]))
83 
84     return [[_execute_task(dsk[k]) for k in keys_d] for keys_d in keys]
85
```

Path: `mars/contrib/dask/converter.py`
Content:
```
1 # Copyright 1999-2021 Alibaba Group Holding Ltd.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #      http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 from dask import is_dask_collection, optimize
16 
17 from .scheduler import mars_dask_get
18 from .utils import reduce
19 
20 
21 def convert_dask_collection(dc):
22     """
23     Convert dask collection object into mars.core.Object via remote API
24 
25     Parameters
26     ----------
27     dc: dask collection
28         Dask collection object to be converted.
29 
30     Returns
31     -------
32     Object
33         Mars Object.
34     """
35     if not is_dask_collection(dc):
36         raise TypeError(f"'{type(dc).__name__}' object is not a valid dask collection")
37 
38     dc.__dask_graph__().validate()
39     dsk = optimize(dc)[0].__dask_graph__()
40 
41     first_key = next(iter(dsk.keys()))
42     if isinstance(first_key, str):
43         key = [first_key]
44     elif isinstance(first_key, tuple):
45         key = sorted([i for i in dsk.keys() if i[0] == first_key[0]], key=lambda x: x[1])
46     else:
47         raise ValueError(
48             f"Dask collection object seems be broken, with unexpected key type:'{type(first_key).__name__}'")
49 
50     return reduce(mars_dask_get(dsk, [key]))
51
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/mars/contrib/dask/converter.py b/mars/contrib/dask/converter.py
--- a/mars/contrib/dask/converter.py
+++ b/mars/contrib/dask/converter.py
@@ -13,9 +13,11 @@
 # limitations under the License.
 
 from dask import is_dask_collection, optimize
+from dask.bag import Bag
 
 from .scheduler import mars_dask_get
 from .utils import reduce
+from ...remote import spawn
 
 
 def convert_dask_collection(dc):
@@ -46,5 +48,8 @@
     else:
         raise ValueError(
             f"Dask collection object seems be broken, with unexpected key type:'{type(first_key).__name__}'")
-
-    return reduce(mars_dask_get(dsk, [key]))
+    res = reduce(mars_dask_get(dsk, [key]))
+    if isinstance(dc, Bag):
+        return spawn(lambda x: list(x[0][0]), args=(res,))
+    else:
+        return res
diff --git a/mars/contrib/dask/scheduler.py b/mars/contrib/dask/scheduler.py
--- a/mars/contrib/dask/scheduler.py
+++ b/mars/contrib/dask/scheduler.py
@@ -12,10 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import List, Tuple
-
 from dask.core import istask, ishashable
 
+from typing import List, Tuple
 from .utils import reduce
 from ...remote import spawn
 
@@ -39,8 +38,11 @@
     Object
         Computed values corresponding to the provided keys.
     """
-
-    return [[reduce(mars_dask_get(dsk, keys)).execute().fetch()]]
+    res = reduce(mars_dask_get(dsk, keys)).execute().fetch()
+    if not isinstance(res, List):
+        return [[res]]
+    else:
+        return res
 
 
 def mars_dask_get(dsk: dict, keys: List[List]):
diff --git a/mars/contrib/dask/utils.py b/mars/contrib/dask/utils.py
--- a/mars/contrib/dask/utils.py
+++ b/mars/contrib/dask/utils.py
@@ -44,8 +44,6 @@
         res = df_concat(objs)
     else:
         res = objs
-        while isinstance(res, List):
-            res = res[0]
     return res.compute() if is_dask_collection(res) else res

{"golden_diff": "diff --git a/mars/contrib/dask/converter.py b/mars/contrib/dask/converter.py\n--- a/mars/contrib/dask/converter.py\n+++ b/mars/contrib/dask/converter.py\n@@ -13,9 +13,11 @@\n # limitations under the License.\n \n from dask import is_dask_collection, optimize\n+from dask.bag import Bag\n \n from .scheduler import mars_dask_get\n from .utils import reduce\n+from ...remote import spawn\n \n \n def convert_dask_collection(dc):\n@@ -46,5 +48,8 @@\n     else:\n         raise ValueError(\n             f\"Dask collection object seems be broken, with unexpected key type:'{type(first_key).__name__}'\")\n-\n-    return reduce(mars_dask_get(dsk, [key]))\n+    res = reduce(mars_dask_get(dsk, [key]))\n+    if isinstance(dc, Bag):\n+        return spawn(lambda x: list(x[0][0]), args=(res,))\n+    else:\n+        return res\ndiff --git a/mars/contrib/dask/scheduler.py b/mars/contrib/dask/scheduler.py\n--- a/mars/contrib/dask/scheduler.py\n+++ b/mars/contrib/dask/scheduler.py\n@@ -12,10 +12,9 @@\n # See the License for the specific language governing permissions and\n # limitations under the License.\n \n-from typing import List, Tuple\n-\n from dask.core import istask, ishashable\n \n+from typing import List, Tuple\n from .utils import reduce\n from ...remote import spawn\n \n@@ -39,8 +38,11 @@\n     Object\n         Computed values corresponding to the provided keys.\n     \"\"\"\n-\n-    return [[reduce(mars_dask_get(dsk, keys)).execute().fetch()]]\n+    res = reduce(mars_dask_get(dsk, keys)).execute().fetch()\n+    if not isinstance(res, List):\n+        return [[res]]\n+    else:\n+        return res\n \n \n def mars_dask_get(dsk: dict, keys: List[List]):\ndiff --git a/mars/contrib/dask/utils.py b/mars/contrib/dask/utils.py\n--- a/mars/contrib/dask/utils.py\n+++ b/mars/contrib/dask/utils.py\n@@ -44,8 +44,6 @@\n         res = df_concat(objs)\n     else:\n         res = objs\n-        while isinstance(res, List):\n-            res = res[0]\n     return res.compute() if is_dask_collection(res) else res\n", "issue": "[BUG] DASK-on-Mars failed when compute multiple objects\n<!--\r\nThank you for your contribution!\r\n\r\nPlease review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.\r\n-->\r\n\r\n**Describe the bug**\r\n\r\nDASK-on-Mars failed when compute multiple objects.\r\n\r\n**To Reproduce**\r\nTo help us reproducing this bug, please provide information below:\r\n1. Your Python version\r\n2. The version of Mars you use\r\n3. Versions of crucial packages, such as numpy, scipy and pandas\r\n4. Full stack of the error.\r\n5. Minimized code to reproduce the error.\r\n\r\n```python\r\nimport dask\r\nimport mars\r\nfrom mars.contrib.dask import mars_scheduler\r\n\r\ndef add(x: int):\r\n    return x + 1\r\n\r\n\r\nmars.new_session()\r\nxs = [dask.delayed(add)(i) for i in range(10)]\r\ndask.compute(xs, scheduler=mars_scheduler)\r\n```\r\n\r\nOutput:\r\n\r\n```\r\n---------------------------------------------------------------------------\r\nIndexError                                Traceback (most recent call last)\r\n<ipython-input-12-531953a29902> in <module>\r\n      4 \r\n      5 xs = [dask.delayed(add)(i) for i in range(10)]\r\n----> 6 dask.compute(xs, scheduler=mars_scheduler)\r\n\r\n~/miniconda3/lib/python3.7/site-packages/dask/base.py in compute(*args, **kwargs)\r\n    567 \r\n    568     results = schedule(dsk, keys, **kwargs)\r\n--> 569     return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])\r\n    570 \r\n    571 \r\n\r\n~/miniconda3/lib/python3.7/site-packages/dask/base.py in repack(results)\r\n    448         dsk = repack_dsk.copy()\r\n    449         dsk[collections_token] = quote(results)\r\n--> 450         return simple_get(dsk, out)\r\n    451 \r\n    452     return collections, repack\r\n\r\n~/miniconda3/lib/python3.7/site-packages/dask/core.py in get(dsk, out, cache)\r\n    149     for key in toposort(dsk):\r\n    150         task = dsk[key]\r\n--> 151         result = _execute_task(task, cache)\r\n    152         cache[key] = result\r\n    153     result = _execute_task(out, cache)\r\n\r\n~/miniconda3/lib/python3.7/site-packages/dask/core.py in _execute_task(arg, cache, dsk)\r\n    119         # temporaries by their reference count and can execute certain\r\n    120         # operations in-place.\r\n--> 121         return func(*(_execute_task(a, cache) for a in args))\r\n    122     elif not ishashable(arg):\r\n    123         return arg\r\n\r\nIndexError: list index out of range\r\n```\r\n\n[BUG] DASK-on-Mars failed when compute multiple objects\n<!--\r\nThank you for your contribution!\r\n\r\nPlease review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.\r\n-->\r\n\r\n**Describe the bug**\r\n\r\nDASK-on-Mars failed when compute multiple objects.\r\n\r\n**To Reproduce**\r\nTo help us reproducing this bug, please provide information below:\r\n1. Your Python version\r\n2. The version of Mars you use\r\n3. Versions of crucial packages, such as numpy, scipy and pandas\r\n4. Full stack of the error.\r\n5. Minimized code to reproduce the error.\r\n\r\n```python\r\nimport dask\r\nimport mars\r\nfrom mars.contrib.dask import mars_scheduler\r\n\r\ndef add(x: int):\r\n    return x + 1\r\n\r\n\r\nmars.new_session()\r\nxs = [dask.delayed(add)(i) for i in range(10)]\r\ndask.compute(xs, scheduler=mars_scheduler)\r\n```\r\n\r\nOutput:\r\n\r\n```\r\n---------------------------------------------------------------------------\r\nIndexError                                Traceback (most recent call last)\r\n<ipython-input-12-531953a29902> in <module>\r\n      4 \r\n      5 xs = [dask.delayed(add)(i) for i in range(10)]\r\n----> 6 dask.compute(xs, scheduler=mars_scheduler)\r\n\r\n~/miniconda3/lib/python3.7/site-packages/dask/base.py in compute(*args, **kwargs)\r\n    567 \r\n    568     results = schedule(dsk, keys, **kwargs)\r\n--> 569     return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])\r\n    570 \r\n    571 \r\n\r\n~/miniconda3/lib/python3.7/site-packages/dask/base.py in repack(results)\r\n    448         dsk = repack_dsk.copy()\r\n    449         dsk[collections_token] = quote(results)\r\n--> 450         return simple_get(dsk, out)\r\n    451 \r\n    452     return collections, repack\r\n\r\n~/miniconda3/lib/python3.7/site-packages/dask/core.py in get(dsk, out, cache)\r\n    149     for key in toposort(dsk):\r\n    150         task = dsk[key]\r\n--> 151         result = _execute_task(task, cache)\r\n    152         cache[key] = result\r\n    153     result = _execute_task(out, cache)\r\n\r\n~/miniconda3/lib/python3.7/site-packages/dask/core.py in _execute_task(arg, cache, dsk)\r\n    119         # temporaries by their reference count and can execute certain\r\n    120         # operations in-place.\r\n--> 121         return func(*(_execute_task(a, cache) for a in args))\r\n    122     elif not ishashable(arg):\r\n    123         return arg\r\n\r\nIndexError: list index out of range\r\n```\r\n\n", "before_files": [{"content": "# Copyright 1999-2021 Alibaba Group Holding Ltd.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#      http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom typing import List\n\nfrom dask import is_dask_collection\nfrom dask.array.core import _concatenate2 as array_concat\nfrom dask.dataframe import concat as df_concat\nfrom dask.utils import is_arraylike, is_dataframe_like, is_series_like, is_index_like\n\nfrom ...remote import spawn\n\n\ndef concat(objs: List):\n    \"\"\"\n    Concat the results of partitioned dask task executions. This function guess the\n        types of resulting list, then calls the corresponding native dask concat functions.\n\n    Parameters\n    ----------\n    objs: List\n        List of the partitioned dask task execution results, which will be concat.\n\n    Returns\n    -------\n    obj:\n        The concat result\n\n    \"\"\"\n    if is_arraylike(objs[0]):\n        res = array_concat(objs, axes=[0])  # TODO: Add concat with args support\n    elif any((is_dataframe_like(objs[0]), is_series_like(objs[0]), is_index_like(objs[0]))):\n        res = df_concat(objs)\n    else:\n        res = objs\n        while isinstance(res, List):\n            res = res[0]\n    return res.compute() if is_dask_collection(res) else res\n\n\ndef reduce(objs: List[List]):\n    \"\"\"\n    Spawn a concat task for 2d-list objects\n\n    Parameters\n    ----------\n    objs: List\n        2d-list of the partitioned dask task execution results, which will be concat.\n\n    Returns\n    -------\n    obj:\n        The spawning concat task.\n    \"\"\"\n    return spawn(\n        concat,\n        args=([spawn(concat, args=(objs_d,)) for objs_d in objs],),\n    )\n", "path": "mars/contrib/dask/utils.py"}, {"content": "# Copyright 1999-2021 Alibaba Group Holding Ltd.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#      http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom typing import List, Tuple\n\nfrom dask.core import istask, ishashable\n\nfrom .utils import reduce\nfrom ...remote import spawn\n\n\ndef mars_scheduler(dsk: dict, keys: List[List[str]]):\n    \"\"\"\n    A Dask-Mars scheduler\n\n    This scheduler is intended to be compatible with existing\n    dask user interface, no callbacks are implemented.\n\n    Parameters\n    ----------\n    dsk: Dict\n        Dask graph, represented as a task DAG dictionary.\n    keys: List[List[str]]\n        2d-list of Dask graph keys whose values we wish to compute and return.\n\n    Returns\n    -------\n    Object\n        Computed values corresponding to the provided keys.\n    \"\"\"\n\n    return [[reduce(mars_dask_get(dsk, keys)).execute().fetch()]]\n\n\ndef mars_dask_get(dsk: dict, keys: List[List]):\n    \"\"\"\n    A Dask-Mars convert function. This function will send the dask graph layers\n        to Mars Remote API, generating mars objects correspond to the provided keys.\n\n    Parameters\n    ----------\n    dsk: Dict\n        Dask graph, represented as a task DAG dictionary.\n    keys: List[List[str]]\n        2d-list of Dask graph keys whose values we wish to compute and return.\n\n    Returns\n    -------\n    Object\n        Spawned mars objects corresponding to the provided keys.\n    \"\"\"\n\n    def _get_arg(a):\n        # if arg contains layer index or callable objs, handle it\n        if ishashable(a) and a in dsk.keys():\n            while ishashable(a) and a in dsk.keys():\n                a = dsk[a]\n            return _execute_task(a)\n        elif not isinstance(a, str) and hasattr(a, \"__getitem__\"):\n            if istask(a):  # TODO:Handle `SubgraphCallable`, which may contains dsk in it\n                return spawn(a[0], args=tuple(_get_arg(i) for i in a[1:]))\n            elif isinstance(a, dict):\n                return {k: _get_arg(v) for k, v in a.items()}\n            elif isinstance(a, List) or isinstance(a, Tuple):\n                return type(a)(_get_arg(i) for i in a)\n        return a\n\n    def _execute_task(task: tuple):\n        if not istask(task):\n            return _get_arg(task)\n        return spawn(task[0], args=tuple(_get_arg(a) for a in task[1:]))\n\n    return [[_execute_task(dsk[k]) for k in keys_d] for keys_d in keys]\n", "path": "mars/contrib/dask/scheduler.py"}, {"content": "# Copyright 1999-2021 Alibaba Group Holding Ltd.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#      http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom dask import is_dask_collection, optimize\n\nfrom .scheduler import mars_dask_get\nfrom .utils import reduce\n\n\ndef convert_dask_collection(dc):\n    \"\"\"\n    Convert dask collection object into mars.core.Object via remote API\n\n    Parameters\n    ----------\n    dc: dask collection\n        Dask collection object to be converted.\n\n    Returns\n    -------\n    Object\n        Mars Object.\n    \"\"\"\n    if not is_dask_collection(dc):\n        raise TypeError(f\"'{type(dc).__name__}' object is not a valid dask collection\")\n\n    dc.__dask_graph__().validate()\n    dsk = optimize(dc)[0].__dask_graph__()\n\n    first_key = next(iter(dsk.keys()))\n    if isinstance(first_key, str):\n        key = [first_key]\n    elif isinstance(first_key, tuple):\n        key = sorted([i for i in dsk.keys() if i[0] == first_key[0]], key=lambda x: x[1])\n    else:\n        raise ValueError(\n            f\"Dask collection object seems be broken, with unexpected key type:'{type(first_key).__name__}'\")\n\n    return reduce(mars_dask_get(dsk, [key]))\n", "path": "mars/contrib/dask/converter.py"}], "after_files": [{"content": "# Copyright 1999-2021 Alibaba Group Holding Ltd.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#      http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom typing import List\n\nfrom dask import is_dask_collection\nfrom dask.array.core import _concatenate2 as array_concat\nfrom dask.dataframe import concat as df_concat\nfrom dask.utils import is_arraylike, is_dataframe_like, is_series_like, is_index_like\n\nfrom ...remote import spawn\n\n\ndef concat(objs: List):\n    \"\"\"\n    Concat the results of partitioned dask task executions. This function guess the\n        types of resulting list, then calls the corresponding native dask concat functions.\n\n    Parameters\n    ----------\n    objs: List\n        List of the partitioned dask task execution results, which will be concat.\n\n    Returns\n    -------\n    obj:\n        The concat result\n\n    \"\"\"\n    if is_arraylike(objs[0]):\n        res = array_concat(objs, axes=[0])  # TODO: Add concat with args support\n    elif any((is_dataframe_like(objs[0]), is_series_like(objs[0]), is_index_like(objs[0]))):\n        res = df_concat(objs)\n    else:\n        res = objs\n    return res.compute() if is_dask_collection(res) else res\n\n\ndef reduce(objs: List[List]):\n    \"\"\"\n    Spawn a concat task for 2d-list objects\n\n    Parameters\n    ----------\n    objs: List\n        2d-list of the partitioned dask task execution results, which will be concat.\n\n    Returns\n    -------\n    obj:\n        The spawning concat task.\n    \"\"\"\n    return spawn(\n        concat,\n        args=([spawn(concat, args=(objs_d,)) for objs_d in objs],),\n    )\n", "path": "mars/contrib/dask/utils.py"}, {"content": "# Copyright 1999-2021 Alibaba Group Holding Ltd.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#      http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom dask.core import istask, ishashable\n\nfrom typing import List, Tuple\nfrom .utils import reduce\nfrom ...remote import spawn\n\n\ndef mars_scheduler(dsk: dict, keys: List[List[str]]):\n    \"\"\"\n    A Dask-Mars scheduler\n\n    This scheduler is intended to be compatible with existing\n    dask user interface, no callbacks are implemented.\n\n    Parameters\n    ----------\n    dsk: Dict\n        Dask graph, represented as a task DAG dictionary.\n    keys: List[List[str]]\n        2d-list of Dask graph keys whose values we wish to compute and return.\n\n    Returns\n    -------\n    Object\n        Computed values corresponding to the provided keys.\n    \"\"\"\n    res = reduce(mars_dask_get(dsk, keys)).execute().fetch()\n    if not isinstance(res, List):\n        return [[res]]\n    else:\n        return res\n\n\ndef mars_dask_get(dsk: dict, keys: List[List]):\n    \"\"\"\n    A Dask-Mars convert function. This function will send the dask graph layers\n        to Mars Remote API, generating mars objects correspond to the provided keys.\n\n    Parameters\n    ----------\n    dsk: Dict\n        Dask graph, represented as a task DAG dictionary.\n    keys: List[List[str]]\n        2d-list of Dask graph keys whose values we wish to compute and return.\n\n    Returns\n    -------\n    Object\n        Spawned mars objects corresponding to the provided keys.\n    \"\"\"\n\n    def _get_arg(a):\n        # if arg contains layer index or callable objs, handle it\n        if ishashable(a) and a in dsk.keys():\n            while ishashable(a) and a in dsk.keys():\n                a = dsk[a]\n            return _execute_task(a)\n        elif not isinstance(a, str) and hasattr(a, \"__getitem__\"):\n            if istask(a):  # TODO:Handle `SubgraphCallable`, which may contains dsk in it\n                return spawn(a[0], args=tuple(_get_arg(i) for i in a[1:]))\n            elif isinstance(a, dict):\n                return {k: _get_arg(v) for k, v in a.items()}\n            elif isinstance(a, List) or isinstance(a, Tuple):\n                return type(a)(_get_arg(i) for i in a)\n        return a\n\n    def _execute_task(task: tuple):\n        if not istask(task):\n            return _get_arg(task)\n        return spawn(task[0], args=tuple(_get_arg(a) for a in task[1:]))\n\n    return [[_execute_task(dsk[k]) for k in keys_d] for keys_d in keys]\n", "path": "mars/contrib/dask/scheduler.py"}, {"content": "# Copyright 1999-2021 Alibaba Group Holding Ltd.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#      http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom dask import is_dask_collection, optimize\nfrom dask.bag import Bag\n\nfrom .scheduler import mars_dask_get\nfrom .utils import reduce\nfrom ...remote import spawn\n\n\ndef convert_dask_collection(dc):\n    \"\"\"\n    Convert dask collection object into mars.core.Object via remote API\n\n    Parameters\n    ----------\n    dc: dask collection\n        Dask collection object to be converted.\n\n    Returns\n    -------\n    Object\n        Mars Object.\n    \"\"\"\n    if not is_dask_collection(dc):\n        raise TypeError(f\"'{type(dc).__name__}' object is not a valid dask collection\")\n\n    dc.__dask_graph__().validate()\n    dsk = optimize(dc)[0].__dask_graph__()\n\n    first_key = next(iter(dsk.keys()))\n    if isinstance(first_key, str):\n        key = [first_key]\n    elif isinstance(first_key, tuple):\n        key = sorted([i for i in dsk.keys() if i[0] == first_key[0]], key=lambda x: x[1])\n    else:\n        raise ValueError(\n            f\"Dask collection object seems be broken, with unexpected key type:'{type(first_key).__name__}'\")\n    res = reduce(mars_dask_get(dsk, [key]))\n    if isinstance(dc, Bag):\n        return spawn(lambda x: list(x[0][0]), args=(res,))\n    else:\n        return res\n", "path": "mars/contrib/dask/converter.py"}]}