Repository: mars-project/mars

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
```



Discussion:
@loopyme Could you please check what went wrong?
Work in progress
@loopyme Could you please check what went wrong?
Work in progress

Resolving patch:
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 @@ def convert_dask_collection(dc):
     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 @@ def mars_scheduler(dsk: dict, keys: List[List[str]]):
     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 @@ def concat(objs: List):
         res = df_concat(objs)
     else:
         res = objs
-        while isinstance(res, List):
-            res = res[0]
     return res.compute() if is_dask_collection(res) else res
 
 


Test patch:
diff --git a/mars/contrib/dask/tests/test_dask.py b/mars/contrib/dask/tests/test_dask.py
--- a/mars/contrib/dask/tests/test_dask.py
+++ b/mars/contrib/dask/tests/test_dask.py
@@ -112,12 +112,25 @@ def test_bag(setup_cluster):
 
     dask_res = result.compute()
     assert dask_res == result.compute(scheduler=mars_scheduler)
-    assert dask_res == list(
-        convert_dask_collection(result).execute().fetch()
-    )  # TODO: dask-bag computation will return weird tuple, which we don't know why
+    assert dask_res == convert_dask_collection(result).execute().fetch()
 
 
 @pytest.mark.skipif(not dask_installed, reason='dask not installed')
 def test_dask_errors():
     with pytest.raises(TypeError):
         convert_dask_collection({"foo": 0, "bar": 1})
+
+
+@pytest.mark.skipif(not dask_installed, reason='dask not installed')
+def test_multiple_objects(setup_cluster):
+    import dask
+
+    def inc(x: int):
+        return x + 1
+
+    test_list = [dask.delayed(inc)(i) for i in range(10)]
+    test_tuple = tuple(dask.delayed(inc)(i) for i in range(10))
+    test_dict = {str(i): dask.delayed(inc)(i) for i in range(10)}
+
+    for test_obj in (test_list, test_tuple, test_dict):
+        assert dask.compute(test_obj) == dask.compute(test_obj, scheduler=mars_scheduler)