We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[BUG] `dask.persist` cannot work on dask-on-mars
<!--
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.persist` cannot work on dask-on-mars.

**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.

```
In [1]: import mars

In [2]: mars.new_session()
Web service started at http://0.0.0.0:64573
Out[2]: <mars.deploy.oscar.session.SyncSession at 0x7fa0f8366d90>

In [3]: import dask
   ...: from mars.contrib.dask import mars_scheduler

In [4]: def inc(x):
   ...:     return x + 1
   ...: 

In [5]: dask_task = dask.delayed(inc)(1)

In [6]: dask_task.persist(scheduler=mars_scheduler)
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-6-095d61a2acc0> in <module>
----> 1 dask_task.persist(scheduler=mars_scheduler)

~/miniconda3/envs/mars3.8/lib/python3.8/site-packages/dask/base.py in persist(self, **kwargs)
    259         dask.base.persist
    260         """
--> 261         (result,) = persist(self, traverse=False, **kwargs)
    262         return result
    263 

~/miniconda3/envs/mars3.8/lib/python3.8/site-packages/dask/base.py in persist(traverse, optimize_graph, scheduler, *args, **kwargs)
    831         postpersists.append((rebuild, a_keys, state))
    832 
--> 833     results = schedule(dsk, keys, **kwargs)
    834     d = dict(zip(keys, results))
    835     results2 = [r({k: d[k] for k in ks}, *s) for r, ks, s in postpersists]

~/Workspace/mars/mars/contrib/dask/scheduler.py in mars_scheduler(dsk, keys)
     39         Computed values corresponding to the provided keys.
     40     """
---> 41     res = reduce(mars_dask_get(dsk, keys)).execute().fetch()
     42     if not isinstance(res, List):
     43         return [[res]]

~/Workspace/mars/mars/contrib/dask/scheduler.py in mars_dask_get(dsk, keys)
     86         return spawn(task[0], args=tuple(_get_arg(a) for a in task[1:]))
     87 
---> 88     return [[_execute_task(dsk[k]) for k in keys_d] for keys_d in keys]

~/Workspace/mars/mars/contrib/dask/scheduler.py in <listcomp>(.0)
     86         return spawn(task[0], args=tuple(_get_arg(a) for a in task[1:]))
     87 
---> 88     return [[_execute_task(dsk[k]) for k in keys_d] for keys_d in keys]

~/Workspace/mars/mars/contrib/dask/scheduler.py in <listcomp>(.0)
     86         return spawn(task[0], args=tuple(_get_arg(a) for a in task[1:]))
     87 
---> 88     return [[_execute_task(dsk[k]) for k in keys_d] for keys_d in keys]

~/miniconda3/envs/mars3.8/lib/python3.8/site-packages/dask/highlevelgraph.py in __getitem__(self, key)
    737                 pass
    738 
--> 739         raise KeyError(key)
    740 
    741     def __len__(self) -> int:

KeyError: 'i'
```




--- 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/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 dask.core import istask, ishashable
16 
17 from typing import List, Tuple
18 from .utils import reduce
19 from ...remote import spawn
20 
21 
22 def mars_scheduler(dsk: dict, keys: List[List[str]]):
23     """
24     A Dask-Mars scheduler
25 
26     This scheduler is intended to be compatible with existing
27     dask user interface, no callbacks are implemented.
28 
29     Parameters
30     ----------
31     dsk: Dict
32         Dask graph, represented as a task DAG dictionary.
33     keys: List[List[str]]
34         2d-list of Dask graph keys whose values we wish to compute and return.
35 
36     Returns
37     -------
38     Object
39         Computed values corresponding to the provided keys.
40     """
41     res = reduce(mars_dask_get(dsk, keys)).execute().fetch()
42     if not isinstance(res, List):
43         return [[res]]
44     else:
45         return res
46 
47 
48 def mars_dask_get(dsk: dict, keys: List[List]):
49     """
50     A Dask-Mars convert function. This function will send the dask graph layers
51         to Mars Remote API, generating mars objects correspond to the provided keys.
52 
53     Parameters
54     ----------
55     dsk: Dict
56         Dask graph, represented as a task DAG dictionary.
57     keys: List[List[str]]
58         2d-list of Dask graph keys whose values we wish to compute and return.
59 
60     Returns
61     -------
62     Object
63         Spawned mars objects corresponding to the provided keys.
64     """
65 
66     def _get_arg(a):
67         # if arg contains layer index or callable objs, handle it
68         if ishashable(a) and a in dsk.keys():
69             while ishashable(a) and a in dsk.keys():
70                 a = dsk[a]
71             return _execute_task(a)
72         elif not isinstance(a, str) and hasattr(a, "__getitem__"):
73             if istask(
74                 a
75             ):  # TODO:Handle `SubgraphCallable`, which may contains dsk in it
76                 return spawn(a[0], args=tuple(_get_arg(i) for i in a[1:]))
77             elif isinstance(a, dict):
78                 return {k: _get_arg(v) for k, v in a.items()}
79             elif isinstance(a, List) or isinstance(a, Tuple):
80                 return type(a)(_get_arg(i) for i in a)
81         return a
82 
83     def _execute_task(task: tuple):
84         if not istask(task):
85             return _get_arg(task)
86         return spawn(task[0], args=tuple(_get_arg(a) for a in task[1:]))
87 
88     return [[_execute_task(dsk[k]) for k in keys_d] for keys_d in keys]
89
```
--- 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/scheduler.py b/mars/contrib/dask/scheduler.py
--- a/mars/contrib/dask/scheduler.py
+++ b/mars/contrib/dask/scheduler.py
@@ -14,12 +14,13 @@
 
 from dask.core import istask, ishashable
 
-from typing import List, Tuple
+from typing import List, Tuple, Union
 from .utils import reduce
 from ...remote import spawn
+from ...deploy.oscar.session import execute
 
 
-def mars_scheduler(dsk: dict, keys: List[List[str]]):
+def mars_scheduler(dsk: dict, keys: Union[List[List[str]], List[str]]):
     """
     A Dask-Mars scheduler
 
@@ -30,22 +31,29 @@
     ----------
     dsk: Dict
         Dask graph, represented as a task DAG dictionary.
-    keys: List[List[str]]
-        2d-list of Dask graph keys whose values we wish to compute and return.
+    keys: Union[List[List[str]], List[str]]
+        1d or 2d list of Dask graph keys whose values we wish to compute and return.
 
     Returns
     -------
     Object
-        Computed values corresponding to the provided keys.
+        Computed values corresponding to the provided keys with same dimension.
     """
-    res = reduce(mars_dask_get(dsk, keys)).execute().fetch()
-    if not isinstance(res, List):
-        return [[res]]
-    else:
-        return res
 
+    if isinstance(keys, List) and not isinstance(keys[0], List):  # 1d keys
+        task = execute(mars_dask_get(dsk, keys))
+        if not isinstance(task, List):
+            task = [task]
+        return map(lambda x: x.fetch(), task)
+    else:  # 2d keys
+        res = execute(reduce(mars_dask_get(dsk, keys))).fetch()
+        if not isinstance(res, List):
+            return [[res]]
+        else:
+            return res
 
-def mars_dask_get(dsk: dict, keys: List[List]):
+
+def mars_dask_get(dsk: dict, keys: Union[List[List[str]], List[str]]):
     """
     A Dask-Mars convert function. This function will send the dask graph layers
         to Mars Remote API, generating mars objects correspond to the provided keys.
@@ -54,13 +62,13 @@
     ----------
     dsk: Dict
         Dask graph, represented as a task DAG dictionary.
-    keys: List[List[str]]
-        2d-list of Dask graph keys whose values we wish to compute and return.
+    keys: Union[List[List[str]], List[str]]
+        1d or 2d list of Dask graph keys whose values we wish to compute and return.
 
     Returns
     -------
     Object
-        Spawned mars objects corresponding to the provided keys.
+        Spawned mars objects corresponding to the provided keys with same dimension.
     """
 
     def _get_arg(a):
@@ -68,7 +76,7 @@
         if ishashable(a) and a in dsk.keys():
             while ishashable(a) and a in dsk.keys():
                 a = dsk[a]
-            return _execute_task(a)
+            return _spawn_task(a)
         elif not isinstance(a, str) and hasattr(a, "__getitem__"):
             if istask(
                 a
@@ -80,9 +88,14 @@
                 return type(a)(_get_arg(i) for i in a)
         return a
 
-    def _execute_task(task: tuple):
+    def _spawn_task(task: tuple):
         if not istask(task):
             return _get_arg(task)
         return spawn(task[0], args=tuple(_get_arg(a) for a in task[1:]))
 
-    return [[_execute_task(dsk[k]) for k in keys_d] for keys_d in keys]
+    return [
+        [_spawn_task(dsk[k]) for k in keys_d]
+        if isinstance(keys_d, List)
+        else _spawn_task(dsk[keys_d])
+        for keys_d in keys
+    ]

{"golden_diff": "diff --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@@ -14,12 +14,13 @@\n \n from dask.core import istask, ishashable\n \n-from typing import List, Tuple\n+from typing import List, Tuple, Union\n from .utils import reduce\n from ...remote import spawn\n+from ...deploy.oscar.session import execute\n \n \n-def mars_scheduler(dsk: dict, keys: List[List[str]]):\n+def mars_scheduler(dsk: dict, keys: Union[List[List[str]], List[str]]):\n     \"\"\"\n     A Dask-Mars scheduler\n \n@@ -30,22 +31,29 @@\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+    keys: Union[List[List[str]], List[str]]\n+        1d or 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+        Computed values corresponding to the provided keys with same dimension.\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+    if isinstance(keys, List) and not isinstance(keys[0], List):  # 1d keys\n+        task = execute(mars_dask_get(dsk, keys))\n+        if not isinstance(task, List):\n+            task = [task]\n+        return map(lambda x: x.fetch(), task)\n+    else:  # 2d keys\n+        res = execute(reduce(mars_dask_get(dsk, keys))).fetch()\n+        if not isinstance(res, List):\n+            return [[res]]\n+        else:\n+            return res\n \n-def mars_dask_get(dsk: dict, keys: List[List]):\n+\n+def mars_dask_get(dsk: dict, keys: Union[List[List[str]], List[str]]):\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@@ -54,13 +62,13 @@\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+    keys: Union[List[List[str]], List[str]]\n+        1d or 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+        Spawned mars objects corresponding to the provided keys with same dimension.\n     \"\"\"\n \n     def _get_arg(a):\n@@ -68,7 +76,7 @@\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+            return _spawn_task(a)\n         elif not isinstance(a, str) and hasattr(a, \"__getitem__\"):\n             if istask(\n                 a\n@@ -80,9 +88,14 @@\n                 return type(a)(_get_arg(i) for i in a)\n         return a\n \n-    def _execute_task(task: tuple):\n+    def _spawn_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+    return [\n+        [_spawn_task(dsk[k]) for k in keys_d]\n+        if isinstance(keys_d, List)\n+        else _spawn_task(dsk[keys_d])\n+        for keys_d in keys\n+    ]\n", "issue": "[BUG] `dask.persist` cannot work on dask-on-mars\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\n`dask.persist` cannot work on dask-on-mars.\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```\r\nIn [1]: import mars\r\n\r\nIn [2]: mars.new_session()\r\nWeb service started at http://0.0.0.0:64573\r\nOut[2]: <mars.deploy.oscar.session.SyncSession at 0x7fa0f8366d90>\r\n\r\nIn [3]: import dask\r\n   ...: from mars.contrib.dask import mars_scheduler\r\n\r\nIn [4]: def inc(x):\r\n   ...:     return x + 1\r\n   ...: \r\n\r\nIn [5]: dask_task = dask.delayed(inc)(1)\r\n\r\nIn [6]: dask_task.persist(scheduler=mars_scheduler)\r\n---------------------------------------------------------------------------\r\nKeyError                                  Traceback (most recent call last)\r\n<ipython-input-6-095d61a2acc0> in <module>\r\n----> 1 dask_task.persist(scheduler=mars_scheduler)\r\n\r\n~/miniconda3/envs/mars3.8/lib/python3.8/site-packages/dask/base.py in persist(self, **kwargs)\r\n    259         dask.base.persist\r\n    260         \"\"\"\r\n--> 261         (result,) = persist(self, traverse=False, **kwargs)\r\n    262         return result\r\n    263 \r\n\r\n~/miniconda3/envs/mars3.8/lib/python3.8/site-packages/dask/base.py in persist(traverse, optimize_graph, scheduler, *args, **kwargs)\r\n    831         postpersists.append((rebuild, a_keys, state))\r\n    832 \r\n--> 833     results = schedule(dsk, keys, **kwargs)\r\n    834     d = dict(zip(keys, results))\r\n    835     results2 = [r({k: d[k] for k in ks}, *s) for r, ks, s in postpersists]\r\n\r\n~/Workspace/mars/mars/contrib/dask/scheduler.py in mars_scheduler(dsk, keys)\r\n     39         Computed values corresponding to the provided keys.\r\n     40     \"\"\"\r\n---> 41     res = reduce(mars_dask_get(dsk, keys)).execute().fetch()\r\n     42     if not isinstance(res, List):\r\n     43         return [[res]]\r\n\r\n~/Workspace/mars/mars/contrib/dask/scheduler.py in mars_dask_get(dsk, keys)\r\n     86         return spawn(task[0], args=tuple(_get_arg(a) for a in task[1:]))\r\n     87 \r\n---> 88     return [[_execute_task(dsk[k]) for k in keys_d] for keys_d in keys]\r\n\r\n~/Workspace/mars/mars/contrib/dask/scheduler.py in <listcomp>(.0)\r\n     86         return spawn(task[0], args=tuple(_get_arg(a) for a in task[1:]))\r\n     87 \r\n---> 88     return [[_execute_task(dsk[k]) for k in keys_d] for keys_d in keys]\r\n\r\n~/Workspace/mars/mars/contrib/dask/scheduler.py in <listcomp>(.0)\r\n     86         return spawn(task[0], args=tuple(_get_arg(a) for a in task[1:]))\r\n     87 \r\n---> 88     return [[_execute_task(dsk[k]) for k in keys_d] for keys_d in keys]\r\n\r\n~/miniconda3/envs/mars3.8/lib/python3.8/site-packages/dask/highlevelgraph.py in __getitem__(self, key)\r\n    737                 pass\r\n    738 \r\n--> 739         raise KeyError(key)\r\n    740 \r\n    741     def __len__(self) -> int:\r\n\r\nKeyError: 'i'\r\n```\r\n\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 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(\n                a\n            ):  # 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"}], "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 dask.core import istask, ishashable\n\nfrom typing import List, Tuple, Union\nfrom .utils import reduce\nfrom ...remote import spawn\nfrom ...deploy.oscar.session import execute\n\n\ndef mars_scheduler(dsk: dict, keys: Union[List[List[str]], 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: Union[List[List[str]], List[str]]\n        1d or 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 with same dimension.\n    \"\"\"\n\n    if isinstance(keys, List) and not isinstance(keys[0], List):  # 1d keys\n        task = execute(mars_dask_get(dsk, keys))\n        if not isinstance(task, List):\n            task = [task]\n        return map(lambda x: x.fetch(), task)\n    else:  # 2d keys\n        res = execute(reduce(mars_dask_get(dsk, keys))).fetch()\n        if not isinstance(res, List):\n            return [[res]]\n        else:\n            return res\n\n\ndef mars_dask_get(dsk: dict, keys: Union[List[List[str]], List[str]]):\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: Union[List[List[str]], List[str]]\n        1d or 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 with same dimension.\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 _spawn_task(a)\n        elif not isinstance(a, str) and hasattr(a, \"__getitem__\"):\n            if istask(\n                a\n            ):  # 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 _spawn_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 [\n        [_spawn_task(dsk[k]) for k in keys_d]\n        if isinstance(keys_d, List)\n        else _spawn_task(dsk[keys_d])\n        for keys_d in keys\n    ]\n", "path": "mars/contrib/dask/scheduler.py"}]}