Repository: mars-project/mars

Issue:
[BUG] index out of range when using Mars with XGBOOST 
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**

IndexError: list assignment index out of range

**To Reproduce**
To help us reproducing this bug, please provide information below:
1. Your Python version

Python 3.7.7 on ray docker 1.9

2. The version of Mars you use

0.9.0a2

3. Versions of crucial packages, such as numpy, scipy and pandas

pip install xgboost
pip install "xgboost_ray"
pip install lightgbm

4. Full stack of the error.


(base) ray@eb0b527fa9ea:~/ray/ray$ python main.py 
2022-02-16 12:39:47,496 WARNING ray.py:301 -- Ray is not started, start the local ray cluster by `ray.init`.
2022-02-16 12:39:50,168 INFO services.py:1340 -- View the Ray dashboard at http://127.0.0.1:8265
2022-02-16 12:39:51,580 INFO driver.py:34 -- Setup cluster with {'ray://ray-cluster-1645043987/0': {'CPU': 8}, 'ray://ray-cluster-1645043987/1': {'CPU': 8}, 'ray://ray-cluster-1645043987/2': {'CPU': 8}, 'ray://ray-cluster-1645043987/3': {'CPU': 8}}
2022-02-16 12:39:51,581 INFO driver.py:36 -- Creating placement group ray-cluster-1645043987 with bundles [{'CPU': 8}, {'CPU': 8}, {'CPU': 8}, {'CPU': 8}].
2022-02-16 12:39:51,716 INFO driver.py:50 -- Create placement group success.
2022-02-16 12:39:52,978 INFO ray.py:479 -- Create supervisor on node ray://ray-cluster-1645043987/0/0 succeeds.
2022-02-16 12:39:53,230 INFO ray.py:489 -- Start services on supervisor ray://ray-cluster-1645043987/0/0 succeeds.
2022-02-16 12:40:07,025 INFO ray.py:498 -- Create 4 workers and start services on workers succeeds.
2022-02-16 12:40:07,036 WARNING ray.py:510 -- Web service started at http://0.0.0.0:46910
  0%|                                                                                                                                                     | 0/100 [00:00<?, ?it/s]
Traceback (most recent call last):
  File "main.py", line 69, in <module>
    main()
  File "main.py", line 35, in main
    df_train, df_test = _load_data(n_samples, n_features, n_classes, test_size=0.2)
  File "main.py", line 25, in _load_data
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=test_size, random_state=shuffle_seed)
  File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/learn/model_selection/_split.py", line 145, in train_test_split
    session=session, **(run_kwargs or dict())
  File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/core/entity/executable.py", line 221, in execute
    ret = execute(*self, session=session, **kw)
  File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/deploy/oscar/session.py", line 1779, in execute
    **kwargs,
  File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/deploy/oscar/session.py", line 1574, in execute
    timeout=self._isolated_session.timeout
  File "/home/ray/anaconda3/lib/python3.7/concurrent/futures/_base.py", line 435, in result
    return self.__get_result()
  File "/home/ray/anaconda3/lib/python3.7/concurrent/futures/_base.py", line 384, in __get_result
    raise self._exception
  File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/deploy/oscar/session.py", line 1725, in _execute
    asyncio.shield(execution_info), progress_update_interval
  File "/home/ray/anaconda3/lib/python3.7/asyncio/tasks.py", line 442, in wait_for
    return fut.result()
  File "/home/ray/anaconda3/lib/python3.7/asyncio/tasks.py", line 630, in _wrap_awaitable
    return (yield from awaitable.__await__())
  File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/deploy/oscar/session.py", line 102, in wait
    return await self._aio_task
  File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/deploy/oscar/session.py", line 903, in _run_in_background
    raise task_result.error.with_traceback(task_result.traceback)
  File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/services/task/supervisor/processor.py", line 57, in inner
    return await func(processor, *args, **kwargs)
  File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/services/task/supervisor/processor.py", line 336, in get_next_stage_processor
    chunk_graph = await self._get_next_chunk_graph(self._chunk_graph_iter)
  File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/services/task/supervisor/processor.py", line 266, in _get_next_chunk_graph
    chunk_graph = await fut
  File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/lib/aio/_threads.py", line 36, in to_thread
    return await loop.run_in_executor(None, func_call)
  File "/home/ray/anaconda3/lib/python3.7/concurrent/futures/thread.py", line 57, in run
    result = self.fn(*self.args, **self.kwargs)
  File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/services/task/supervisor/processor.py", line 261, in next_chunk_graph
    return next(chunk_graph_iter)
  File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/services/task/supervisor/preprocessor.py", line 158, in tile
    for chunk_graph in chunk_graph_builder.build():
  File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/core/graph/builder/chunk.py", line 272, in build
    yield from self._build()
  File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/core/graph/builder/chunk.py", line 266, in _build
    graph = next(tile_iterator)
  File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/services/task/supervisor/preprocessor.py", line 75, in __iter__
    to_update_tileables = self._iter()
  File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/core/graph/builder/chunk.py", line 204, in _iter
    visited,
  File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/core/graph/builder/chunk.py", line 113, in _tile
    need_process = next(tile_handler)
  File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/core/graph/builder/chunk.py", line 84, in _tile_handler
    tiled_tileables = yield from handler.tile(tiled_tileables)
  File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/core/entity/tileables.py", line 79, in tile
    tiled_result = yield from tile_handler(op)
  File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/learn/utils/shuffle.py", line 217, in tile
    inp = yield from cls._safe_rechunk(inp, ax_nsplit)
  File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/learn/utils/shuffle.py", line 144, in _safe_rechunk
    return (yield from recursive_tile(tileable.rechunk(ax_nsplit)))
  File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/tensor/rechunk/rechunk.py", line 103, in rechunk
    chunk_size = get_nsplits(tensor, chunk_size, tensor.dtype.itemsize)
  File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/tensor/rechunk/core.py", line 34, in get_nsplits
    chunk_size[idx] = c
IndexError: list assignment index out of range
(RayMainPool pid=5574) Unexpected error happens in <function TaskProcessor.get_next_stage_processor at 0x7f276f2458c0>
(RayMainPool pid=5574) Traceback (most recent call last):
(RayMainPool pid=5574)   File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/services/task/supervisor/processor.py", line 57, in inner
(RayMainPool pid=5574)     return await func(processor, *args, **kwargs)
(RayMainPool pid=5574)   File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/services/task/supervisor/processor.py", line 336, in get_next_stage_processor
(RayMainPool pid=5574)     chunk_graph = await self._get_next_chunk_graph(self._chunk_graph_iter)
(RayMainPool pid=5574)   File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/services/task/supervisor/processor.py", line 266, in _get_next_chunk_graph
(RayMainPool pid=5574)     chunk_graph = await fut
(RayMainPool pid=5574)   File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/lib/aio/_threads.py", line 36, in to_thread
(RayMainPool pid=5574)     return await loop.run_in_executor(None, func_call)
(RayMainPool pid=5574)   File "/home/ray/anaconda3/lib/python3.7/concurrent/futures/thread.py", line 57, in run
(RayMainPool pid=5574)     result = self.fn(*self.args, **self.kwargs)
(RayMainPool pid=5574)   File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/services/task/supervisor/processor.py", line 261, in next_chunk_graph
(RayMainPool pid=5574)     return next(chunk_graph_iter)
(RayMainPool pid=5574)   File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/services/task/supervisor/preprocessor.py", line 158, in tile
(RayMainPool pid=5574)     for chunk_graph in chunk_graph_builder.build():
(RayMainPool pid=5574)   File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/core/graph/builder/chunk.py", line 272, in build
(RayMainPool pid=5574)     yield from self._build()
(RayMainPool pid=5574)   File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/core/graph/builder/chunk.py", line 266, in _build
(RayMainPool pid=5574)     graph = next(tile_iterator)
(RayMainPool pid=5574)   File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/services/task/supervisor/preprocessor.py", line 75, in __iter__
(RayMainPool pid=5574)     to_update_tileables = self._iter()
(RayMainPool pid=5574)   File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/core/graph/builder/chunk.py", line 204, in _iter
(RayMainPool pid=5574)     visited,
(RayMainPool pid=5574)   File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/core/graph/builder/chunk.py", line 113, in _tile
(RayMainPool pid=5574)     need_process = next(tile_handler)
(RayMainPool pid=5574)   File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/core/graph/builder/chunk.py", line 84, in _tile_handler
(RayMainPool pid=5574)     tiled_tileables = yield from handler.tile(tiled_tileables)
(RayMainPool pid=5574)   File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/core/entity/tileables.py", line 79, in tile
(RayMainPool pid=5574)     tiled_result = yield from tile_handler(op)
(RayMainPool pid=5574)   File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/learn/utils/shuffle.py", line 217, in tile
(RayMainPool pid=5574)     inp = yield from cls._safe_rechunk(inp, ax_nsplit)
(RayMainPool pid=5574)   File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/learn/utils/shuffle.py", line 144, in _safe_rechunk
(RayMainPool pid=5574)     return (yield from recursive_tile(tileable.rechunk(ax_nsplit)))
(RayMainPool pid=5574)   File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/tensor/rechunk/rechunk.py", line 103, in rechunk
(RayMainPool pid=5574)     chunk_size = get_nsplits(tensor, chunk_size, tensor.dtype.itemsize)
(RayMainPool pid=5574)   File "/home/ray/anaconda3/lib/python3.7/site-packages/mars/tensor/rechunk/core.py", line 34, in get_nsplits
(RayMainPool pid=5574)     chunk_size[idx] = c
(RayMainPool pid=5574) IndexError: list assignment index out of range



5. Minimized code to reproduce the error.
```python

import logging
import ray
import mars
import numpy as np
import mars.dataframe as md
from mars.learn.model_selection import train_test_split
from mars.learn.datasets import make_classification
from xgboost_ray import RayDMatrix, RayParams, train, predict

logger = logging.getLogger(__name__)
logging.basicConfig(format=ray.ray_constants.LOGGER_FORMAT, level=logging.INFO)

def _load_data(n_samples: int,
               n_features:int,
               n_classes: int,
               test_size: float = 0.1,
               shuffle_seed: int = 42):
    n_informative = int(n_features * 0.5)
    n_redundant = int(n_features * 0.2)
    # generate dataset
    X, y = make_classification(n_samples=n_samples, n_features=n_features, n_classes=n_classes, n_informative=n_informative, n_redundant=n_redundant, random_state=shuffle_seed)
    X, y = md.DataFrame(X), md.DataFrame({"labels": y})
    X.columns = ['feature-' + str(i) for i in range(n_features)]
    # split dataset
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=test_size, random_state=shuffle_seed)
    return md.concat([X_train, y_train], axis=1), md.concat([X_test, y_test], axis=1)

def main(*args):
    n_samples, n_features, worker_num, worker_cpu, num_shards = 10 ** 4, 20, 4, 8, 10
    ray_params = RayParams(num_actors=10, cpus_per_actor=1)

    # setup mars
    mars.new_ray_session(worker_num=worker_num, worker_cpu=worker_cpu, worker_mem=1 * 1024 ** 3)
    n_classes = 10
    df_train, df_test = _load_data(n_samples, n_features, n_classes, test_size=0.2)
    print(df_train)
    print(df_test)
    # convert mars DataFrame to Ray dataset
    ds_train = md.to_ray_dataset(df_train, num_shards=num_shards)
    ds_test = md.to_ray_dataset(df_test, num_shards=num_shards)
    train_set = RayDMatrix(data=ds_train, label="labels")
    test_set = RayDMatrix(data=ds_test, label="labels")

    evals_result = {}
    params = {
        'nthread': 1,
        'objective': 'multi:softmax',
        'eval_metric': ['mlogloss', 'merror'],
        'num_class': n_classes,
        'eta': 0.1,
        'seed': 42
    }
    bst = train(
        params=params,
        dtrain=train_set,
        num_boost_round=200,
        evals=[(train_set, 'train')],
        evals_result=evals_result,
        verbose_eval=100,
        ray_params=ray_params
    )
    # predict on a test set.
    pred = predict(bst, test_set, ray_params=ray_params)
    precision = (ds_test.dataframe['labels'].to_pandas() == pred).astype(int).sum() / ds_test.dataframe.shape[0]
    logger.info("Prediction Accuracy: %.4f", precision)


if __name__ == "__main__":
    main()

```


**Expected behavior**

train_test_split would fail and emit "IndexError: list assignment index out of range".


**Additional context**
Add any other context about the problem here.



Discussion:
I can reproduce the error, will fix it ASAP.

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
@@ -214,7 +214,8 @@ def tile(cls, op):
         ax_nsplit = {ax: decide_unify_split(*ns) for ax, ns in axis_to_nsplits.items()}
         rechunked_inputs = []
         for inp in inputs:
-            inp = yield from cls._safe_rechunk(inp, ax_nsplit)
+            inp_ax_nsplit = {ax: ns for ax, ns in ax_nsplit.items() if ax < inp.ndim}
+            inp = yield from cls._safe_rechunk(inp, inp_ax_nsplit)
             rechunked_inputs.append(inp)
         inputs = rechunked_inputs
 


Test patch:
diff --git a/mars/learn/utils/tests/test_shuffle.py b/mars/learn/utils/tests/test_shuffle.py
--- a/mars/learn/utils/tests/test_shuffle.py
+++ b/mars/learn/utils/tests/test_shuffle.py
@@ -142,18 +142,20 @@ def test_shuffle_execution(setup):
     # tensor + dataframe + series
     raw5 = np.random.rand(10, 15, 20)
     t5 = mt.array(raw5, chunk_size=8)
+    t6 = mt.array(raw5[:, 0, 0], chunk_size=6)
     raw6 = pd.DataFrame(np.random.rand(10, 15))
     df = md.DataFrame(raw6, chunk_size=(8, 15))
     raw7 = pd.Series(np.random.rand(10))
     series = md.Series(raw7, chunk_size=8)
 
     for axes in [(0,), (1,), (0, 1), (1, 2), [0, 1, 2]]:
-        ret = shuffle(t5, df, series, axes=axes, random_state=0)
+        ret = shuffle(t5, df, series, t6, axes=axes, random_state=0)
         # skip check nsplits because it's updated
-        res5, res_df, res_series = ret.execute(
+        res5, res_df, res_series, res6 = ret.execute(
             extra_config={"check_nsplits": False}
         ).fetch(extra_config={"check_nsplits": False})
 
         assert res5.shape == raw5.shape
         assert res_df.shape == df.shape
         assert res_series.shape == series.shape
+        assert res6.shape == (raw5.shape[0],)