# Buggy code
def _update_tileable_and_its_chunk_shapes(self):
    need_update_tileable_to_tiled = dict()
    for tileable in self._chunk_graph_builder.prev_tileable_graph:
        if tileable.key in self._target_tileable_finished:
            tiled = self._tileable_key_opid_to_tiled[tileable.key, tileable.op.id][-1]
            if not has_unknown_shape(tiled):
                continue
            need_update_tileable_to_tiled[tileable] = tiled

    if len(need_update_tileable_to_tiled) == 0:
        return

    need_update_chunks = list(
        c for t in need_update_tileable_to_tiled.values() for c in t.chunks
    )
    chunk_metas = self.chunk_meta.batch_get_chunk_meta(
        self._session_id, list(c.key for c in need_update_chunks)
    )
    for chunk, chunk_meta in zip(need_update_chunks, chunk_metas):
        chunk.data._shape = chunk_meta.chunk_shape

    for tileable, tiled in need_update_tileable_to_tiled.items():
        chunk_idx_to_shape = OrderedDict((c.index, c.shape) for c in tiled.chunks)
        nsplits = calc_nsplits(chunk_idx_to_shape)
        tiled._nsplits = nsplits
        if any(np.isnan(s) for s in tileable.shape):
            shape = tuple(sum(ns) for ns in nsplits)
            tileable._update_shape(shape)
            tiled._update_shape(shape)


# Traceback (KeyError)
2020-12-02 11:19:40,309 mars.scheduler.operands.common 87 ERROR    Attempt 1: Unexpected error KeyError occurred in executing operand 5c7a3b06d448300987640036d2f5a34e in 11.238.145.234:49708
Traceback (most recent call last):
File "/home/admin/work/_public-mars-0.5.5.zip/mars/promise.py", line 378, in _wrapped
return func(*args, **kwargs)
File "/home/admin/work/_public-mars-0.5.5.zip/mars/utils.py", line 377, in _wrapped
return func(*args, **kwargs)
File "/home/admin/work/_public-mars-0.5.5.zip/mars/worker/execution.py", line 564, in execute_graph
quota_request = self._prepare_quota_request(session_id, graph_key)
File "/home/admin/work/_public-mars-0.5.5.zip/mars/worker/execution.py", line 249, in _prepare_quota_request
memory_estimations = self._estimate_calc_memory(session_id, graph_key)
File "/home/admin/work/_public-mars-0.5.5.zip/mars/worker/execution.py", line 213, in _estimate_calc_memory
res = executor.execute_graph(graph_record.graph, graph_record.chunk_targets, mock=True)
File "/home/admin/work/_public-mars-0.5.5.zip/mars/executor.py", line 690, in execute_graph
res = graph_execution.execute(retval)
File "/home/admin/work/_public-mars-0.5.5.zip/mars/executor.py", line 574, in execute
return [self._chunk_results[key] for key in self._keys]
File "/home/admin/work/_public-mars-0.5.5.zip/mars/executor.py", line 574, in <listcomp>
return [self._chunk_results[key] for key in self._keys]
KeyError: '3990ec90331559138b6ecbc6d76fbd0d'

# Fixed code
def _update_tileable_and_its_chunk_shapes(self):
    need_update_tileable_to_tiled = dict()
    for tileable in self._chunk_graph_builder.prev_tileable_graph:
        if tileable.key in self._target_tileable_finished:
            tiled = self._tileable_key_opid_to_tiled[tileable.key, tileable.op.id][-1]
            if not has_unknown_shape(tiled):
                continue
            need_update_tileable_to_tiled[tileable] = tiled

    if len(need_update_tileable_to_tiled) == 0:
        return

    need_update_chunks = list(
        c for t in need_update_tileable_to_tiled.values() for c in t.chunks
    )
    chunk_metas = self.chunk_meta.batch_get_chunk_meta(
        self._session_id, list(c.key for c in need_update_chunks)
    )
    ops_to_restart = set()
    keys = []
    for chunk, chunk_meta in zip(need_update_chunks, chunk_metas):
        if chunk_meta is None:
            ops_to_restart.add(chunk.op.key)
            keys.append(chunk.key)
        else:
            chunk.data._shape = chunk_meta.chunk_shape
    if ops_to_restart:
        for op_key in ops_to_restart:
            self._get_operand_ref(op_key).start_operand(
                OperandState.READY, _tell=True, _wait=False
            )
        raise RuntimeError(
            f"Cannot find chunks {keys}. Operands {ops_to_restart} restarted."
        )

    for tileable, tiled in need_update_tileable_to_tiled.items():
        chunk_idx_to_shape = OrderedDict((c.index, c.shape) for c in tiled.chunks)
        nsplits = calc_nsplits(chunk_idx_to_shape)
        tiled._nsplits = nsplits
        if any(np.isnan(s) for s in tileable.shape):
            shape = tuple(sum(ns) for ns in nsplits)
            tileable._update_shape(shape)
            tiled._update_shape(shape)