# Buggy code
def add_finished_predecessor(
    self, op_key, worker, output_sizes=None, output_shapes=None
):
    super().add_finished_predecessor(
        op_key, worker, output_sizes=output_sizes, output_shapes=output_shapes
    )

    from ..chunkmeta import WorkerMeta

    chunk_key = next(iter(output_sizes.keys()))[0]
    self._mapper_op_to_chunk[op_key] = chunk_key
    if op_key not in self._worker_to_mappers[worker]:
        self._worker_to_mappers[worker].add(op_key)
        self.chunk_meta.add_worker(self._session_id, chunk_key, worker, _tell=True)

    shuffle_keys_to_op = self._shuffle_keys_to_op

    if not self._reducer_workers:
        self._reducer_workers = self._graph_refs[0].assign_operand_workers(
            self._succ_keys, input_chunk_metas=self._reducer_to_mapper
        )
    reducer_workers = self._reducer_workers
    data_to_addresses = dict()

    for (chunk_key, shuffle_key), data_size in output_sizes.items() or ():
        succ_op_key = shuffle_keys_to_op[shuffle_key]
        meta = self._reducer_to_mapper[succ_op_key][op_key] = WorkerMeta(
            chunk_size=data_size,
            workers=(worker,),
            chunk_shape=output_shapes.get((chunk_key, shuffle_key)),
        )
        reducer_worker = reducer_workers.get(succ_op_key)
        if reducer_worker and reducer_worker != worker:
            data_to_addresses[(chunk_key, shuffle_key)] = [reducer_worker]
            meta.workers += (reducer_worker,)

    if data_to_addresses:
        try:
            with rewrite_worker_errors():
                self._get_raw_execution_ref(address=worker).send_data_to_workers(
                    self._session_id, data_to_addresses, _tell=True
                )
        except WorkerDead:
            self._resource_ref.detach_dead_workers([worker], _tell=True)

    if all(k in self._finish_preds for k in self._pred_keys):
        self._start_successors()


# Traceback (KeyError)
Error
Traceback (most recent call last):
File "/Users/wenjun.swj/Code/mars/mars/scheduler/graph.py", line 410, in _execute_graph
self.prepare_graph(compose=compose)
File "/Users/wenjun.swj/Code/mars/mars/utils.py", line 377, in _wrapped
return func(*args, **kwargs)
File "/Users/wenjun.swj/Code/mars/mars/utils.py", line 451, in _inner
return func(*args, **kwargs)
File "/Users/wenjun.swj/Code/mars/mars/scheduler/graph.py", line 646, in prepare_graph
cur_chunk_graph = chunk_graph_builder.build(
File "/Users/wenjun.swj/Code/mars/mars/utils.py", line 451, in _inner
return func(*args, **kwargs)
File "/Users/wenjun.swj/Code/mars/mars/tiles.py", line 347, in build
chunk_graph = super().build(
File "/Users/wenjun.swj/Code/mars/mars/utils.py", line 451, in _inner
return func(*args, **kwargs)
File "/Users/wenjun.swj/Code/mars/mars/tiles.py", line 262, in build
self._on_tile_failure(tileable_data.op, exc_info)
File "/Users/wenjun.swj/Code/mars/mars/tiles.py", line 301, in inner
raise exc_info[1].with_traceback(exc_info[2]) from None
File "/Users/wenjun.swj/Code/mars/mars/tiles.py", line 242, in build
tiled = self._tile(tileable_data, tileable_graph)
File "/Users/wenjun.swj/Code/mars/mars/tiles.py", line 337, in _tile
return super()._tile(tileable_data, tileable_graph)
File "/Users/wenjun.swj/Code/mars/mars/tiles.py", line 203, in _tile
tds = on_tile(tileable_data.op.outputs, tds)
File "/Users/wenjun.swj/Code/mars/mars/scheduler/graph.py", line 630, in on_tile
return self.context.wraps(handler.dispatch)(first.op)
File "/Users/wenjun.swj/Code/mars/mars/context.py", line 72, in h
return func(*args, **kwargs)
File "/Users/wenjun.swj/Code/mars/mars/utils.py", line 451, in _inner
return func(*args, **kwargs)
File "/Users/wenjun.swj/Code/mars/mars/tiles.py", line 119, in dispatch
tiled = op_cls.tile(op)
File "/Users/wenjun.swj/Code/mars/mars/learn/contrib/xgboost/dmatrix.py", line 257, in tile
return cls._tile_single_output(op)
File "/Users/wenjun.swj/Code/mars/mars/learn/contrib/xgboost/dmatrix.py", line 227, in _tile_single_output
data_chunk = concat_chunks(chunks[0])
File "/Users/wenjun.swj/Code/mars/mars/learn/utils/core.py", line 33, in concat_chunks
tileable = chunks[0].op.create_tileable_from_chunks(chunks)
File "/Users/wenjun.swj/Code/mars/mars/dataframe/operands.py", line 212, in create_tileable_from_chunks
params = cls._calc_dataframe_params(chunk_index_to_chunk, chunk_shape)
File "/Users/wenjun.swj/Code/mars/mars/dataframe/operands.py", line 228, in _calc_dataframe_params
pd_indxes = [chunk_index_to_chunks[i, 0].index_value.to_pandas()
File "/Users/wenjun.swj/Code/mars/mars/dataframe/operands.py", line 228, in <listcomp>
pd_indxes = [chunk_index_to_chunks[i, 0].index_value.to_pandas()
KeyError: (2, 0)

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
File "/Users/wenjun.swj/miniconda3/lib/python3.8/unittest/case.py", line 60, in testPartExecutor
yield
File "/Users/wenjun.swj/miniconda3/lib/python3.8/unittest/case.py", line 676, in run
self._callTestMethod(testMethod)
File "/Users/wenjun.swj/miniconda3/lib/python3.8/unittest/case.py", line 633, in _callTestMethod
method()
File "/Users/wenjun.swj/Code/mars/mars/learn/contrib/xgboost/tests/integrated/test_distributed_xgboost.py", line 71, in testDistributedXGBClassifier
classifier.fit(X, y, eval_set=[(X, y)], session=sess, run_kwargs=run_kwargs)
File "/Users/wenjun.swj/Code/mars/mars/learn/contrib/xgboost/classifier.py", line 55, in fit
result = train(params, dtrain, num_boost_round=self.get_num_boosting_rounds(),
File "/Users/wenjun.swj/Code/mars/mars/learn/contrib/xgboost/train.py", line 200, in train
ret = t.execute(session=session, **run_kwargs).fetch(session=session)
File "/Users/wenjun.swj/Code/mars/mars/core.py", line 379, in execute
return run()
File "/Users/wenjun.swj/Code/mars/mars/core.py", line 374, in run
session.run(self, **kw)
File "/Users/wenjun.swj/Code/mars/mars/session.py", line 499, in run
result = self._sess.run(*tileables, **kw)
File "/Users/wenjun.swj/Code/mars/mars/web/session.py", line 214, in run
if self._check_response_finished(graph_url, timeout_val):
File "/Users/wenjun.swj/Code/mars/mars/web/session.py", line 174, in _check_response_finished
raise ExecutionFailed('Graph execution failed.') from exc
mars.errors.ExecutionFailed: 'Graph execution failed.'

# Fixed code
def add_finished_predecessor(
    self, op_key, worker, output_sizes=None, output_shapes=None
):
    super().add_finished_predecessor(
        op_key, worker, output_sizes=output_sizes, output_shapes=output_shapes
    )

    from ..chunkmeta import WorkerMeta

    chunk_key = next(iter(output_sizes.keys()))[0]
    self._mapper_op_to_chunk[op_key] = chunk_key
    if op_key not in self._worker_to_mappers[worker]:
        self._worker_to_mappers[worker].add(op_key)
        self.chunk_meta.add_worker(self._session_id, chunk_key, worker, _tell=True)

    shuffle_keys_to_op = self._shuffle_keys_to_op

    if not self._reducer_workers:
        self._reducer_workers = self._graph_refs[0].assign_operand_workers(
            self._succ_keys, input_chunk_metas=self._reducer_to_mapper
        )
    reducer_workers = self._reducer_workers
    data_to_addresses = dict()

    unused_keys = []

    for (chunk_key, shuffle_key), data_size in output_sizes.items() or ():
        if shuffle_key not in shuffle_keys_to_op:
            # outputs may be pruned, hence those keys become useless
            unused_keys.append((chunk_key, shuffle_key))
            continue

        succ_op_key = shuffle_keys_to_op[shuffle_key]
        meta = self._reducer_to_mapper[succ_op_key][op_key] = WorkerMeta(
            chunk_size=data_size,
            workers=(worker,),
            chunk_shape=output_shapes.get((chunk_key, shuffle_key)),
        )
        reducer_worker = reducer_workers.get(succ_op_key)
        if reducer_worker and reducer_worker != worker:
            data_to_addresses[(chunk_key, shuffle_key)] = [reducer_worker]
            meta.workers += (reducer_worker,)

    if unused_keys:
        self._free_data_in_worker(unused_keys, [(worker,)] * len(unused_keys))

    if data_to_addresses:
        try:
            with rewrite_worker_errors():
                self._get_raw_execution_ref(address=worker).send_data_to_workers(
                    self._session_id, data_to_addresses, _tell=True
                )
        except WorkerDead:
            self._resource_ref.detach_dead_workers([worker], _tell=True)

    if all(k in self._finish_preds for k in self._pred_keys):
        self._start_successors()