# Copyright 1999-2020 Alibaba Group Holding Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import contextlib
import itertools
import logging
import operator
import os
import random
import sys
import time
from collections import defaultdict, OrderedDict
from functools import lru_cache, reduce

import numpy as np

from ..actors.errors import ActorAlreadyExist
from ..config import options
from ..context import DistributedContext
from ..core import get_tiled, TileableGraph, ChunkGraph
from ..core.entity.tileables import handler
from ..core.graph import DAG, IterativeChunkGraphBuilder, TileableGraphBuilder
from ..core.operand import Fetch, ShuffleProxy, VirtualOperand, SuccessorsExclusive
from ..errors import ExecutionInterrupted, GraphNotExists, WorkerDead
from ..optimizes.tileable_graph import OptimizeIntegratedTileableGraphBuilder
from ..serialize import dataserializer
from ..utils import serialize_graph, deserialize_graph, log_unhandled, \
    build_exc_info, build_fetch_chunk, build_fetch_tileable, calc_nsplits, \
    get_chunk_reducer_index, enter_mode, has_unknown_shape
from .analyzer import GraphAnalyzer
from .assigner import AssignerActor
from .kvstore import KVStoreActor
from .operands import get_operand_actor_class, OperandState
from .resource import ResourceActor
from .session import SessionActor
from .utils import SchedulerActor, GraphState

logger = logging.getLogger(__name__)


class GraphMetaActor(SchedulerActor):
    """
    Actor storing metadata of a graph
    """
    @staticmethod
    def gen_uid(session_id, graph_key):
        return f's:0:graph_meta${session_id}${graph_key}'

    def __init__(self, session_id, graph_key, session_ref):
        super().__init__()
        self._session_id = session_id
        self._graph_key = graph_key

        self._session_ref = session_ref
        self._kv_store_ref = None
        self._graph_wait_ref = None

        self._start_time = None
        self._end_time = None
        self._state = None
        self._final_state = None
        self._exc = None

        self._graph_finish_event = None

        self._op_infos = defaultdict(dict)
        self._state_to_infos = defaultdict(dict)

    def post_create(self):
        super().post_create()
        self._kv_store_ref = self.ctx.actor_ref(KVStoreActor.default_uid())
        if not self.ctx.has_actor(self._kv_store_ref):
            self._kv_store_ref = None

        self._session_ref = self.ctx.actor_ref(self._session_ref) if self._session_ref else None
        if not self.ctx.has_actor(self._session_ref):
            self._session_ref = None

        self._graph_finish_event = self.ctx.event()
        graph_wait_uid = GraphWaitActor.gen_uid(self._session_id, self._graph_key)
        try:
            graph_wait_uid = self.ctx.distributor.make_same_process(
                graph_wait_uid, self.uid)
        except AttributeError:
            pass
        self._graph_wait_ref = self.ctx.create_actor(
            GraphWaitActor, self._graph_finish_event, uid=graph_wait_uid)

    def pre_destroy(self):
        self._graph_wait_ref.destroy()
        super().pre_destroy()

    def _update_session_graph_info(self):
        if self._session_ref is not None:
            self._session_ref.set_graph_info(
                self._graph_key, start_time=self._start_time, end_time=self._end_time,
                state=self._state, graph_size=len(self._op_infos), _tell=True, _wait=False)

    def set_graph_start(self):
        self._start_time = time.time()
        self._update_session_graph_info()

    def set_graph_end(self):
        self._end_time = time.time()
        self._graph_finish_event.set()
        self._update_session_graph_info()

    def set_state(self, state):
        self._state = state
        if self._kv_store_ref is not None:
            self._kv_store_ref.write(
                f'/sessions/{self._session_id}/graph/{self._graph_key}/state', state.name,
                _tell=True, _wait=False)
        self._update_session_graph_info()

    def get_state(self):
        return self._state

    def set_exc_info(self, exc):
        self._exc = exc

    def get_exc_info(self):
        return self._exc

    def get_wait_ref(self):
        return self._graph_wait_ref

    def set_final_state(self, state):
        self._final_state = state
        if self._kv_store_ref is not None:
            self._kv_store_ref.write(
                f'/sessions/{self._session_id}/graph/{self._graph_key}/final_state',
                state.name, _tell=True)

    def get_final_state(self):
        return self._final_state

    def get_operand_info(self, state=None):
        return self._op_infos if state is None else self._state_to_infos[state]

    def update_op_state(self, op_key, op_name, op_state):
        new_info = dict(op_name=op_name, state=op_state)
        old_state = self._op_infos[op_key].get('state')
        self._op_infos[op_key].update(new_info)
        if old_state is not None:
            self._state_to_infos[old_state].pop(op_key, None)
        self._state_to_infos[op_state][op_key] = self._op_infos[op_key]

    def update_op_worker(self, op_key, op_name, worker):
        new_info = dict(op_name=op_name, worker=worker)
        self._op_infos[op_key].update(new_info)

    def update_op_infos(self, op_infos):
        for key, info in op_infos.items():
            old_state = self._op_infos[key].get('state')
            if old_state is not None:
                self._state_to_infos[old_state].pop(key, None)
            if info.get('state') is not None:
                self._state_to_infos[info['state']][key] = info

            self._op_infos[key].update(info)
        self._update_session_graph_info()

    @log_unhandled
    def calc_stats(self):
        states = list(OperandState.__members__.values())
        state_mapping = OrderedDict((v, idx) for idx, v in enumerate(states))
        state_names = [s.name for s in state_mapping]

        op_stats = OrderedDict()
        finished = 0
        total_count = len(self._op_infos)
        for operand_info in self._op_infos.values():
            if operand_info.get('virtual'):
                total_count -= 1
                continue
            op_name = operand_info['op_name']
            state = operand_info['state'] or OperandState.UNSCHEDULED
            if state in (OperandState.FINISHED, OperandState.FREED):
                finished += 1
            try:
                stats_list = op_stats[op_name]
            except KeyError:
                stats_list = op_stats[op_name] = [0] * len(state_mapping)
            stats_list[state_mapping[state]] += 1

        data_src = OrderedDict([('states', state_names), ])
        for op, state_stats in op_stats.items():
            sum_chunks = sum(state_stats)
            data_src[op] = [v * 100.0 / sum_chunks for v in state_stats]

        ops = list(data_src)[1:]
        states = data_src['states']
        transposed = OrderedDict()
        transposed['ops'] = ops
        for sid, state in enumerate(states):
            transposed[state] = list()
            for op in ops:
                transposed[state].append(data_src[op][sid])

        percentage = finished * 100.0 / total_count if total_count != 0 else 1
        return ops, transposed, percentage


class GraphWaitActor(SchedulerActor):
    @staticmethod
    def gen_uid(session_id, graph_key):
        return f's:0:graph_wait${session_id}${graph_key}'

    def __init__(self, graph_event):
        super().__init__()
        self._graph_event = graph_event

    def wait(self, timeout=None):
        self._graph_event.wait(timeout)


class GraphActor(SchedulerActor):
    """
    Actor handling execution and status of a Mars graph
    """
    @staticmethod
    def gen_uid(session_id, graph_key):
        return f's:0:graph${session_id}${graph_key}'

    def __init__(self, session_id, graph_key, serialized_tileable_graph,
                 target_tileables=None, serialized_chunk_graph=None,
                 state=GraphState.UNSCHEDULED, final_state=None):
        super().__init__()
        self._graph_key = graph_key
        self._session_id = session_id
        self._serialized_tileable_graph = serialized_tileable_graph
        self._serialized_chunk_graph = serialized_chunk_graph
        self._state = state
        self._final_state = final_state
        self._context = None

        self._operand_free_paused = False

        self._cluster_info_ref = None
        self._assigner_actor_ref = None
        self._resource_actor_ref = None
        self._kv_store_ref = None
        self._graph_meta_ref = None
        self._session_ref = None

        self._tileable_graph_cache = None
        self._chunk_graph_cache = None

        # chunk graph builder
        self._chunk_graph_builder = None

        self._resource_actor = None
        # generated once
        self._target_tileable_keys = target_tileables
        self._target_tileable_datas = []
        self._tileable_key_to_opid = dict()
        # used over all iteration
        self._tileable_key_opid_to_tiled = defaultdict(list)
        self._chunk_key_id_to_chunk = dict()
        self._all_terminated_tileable_keys = set()
        # used only for current iteration
        self._tileable_to_fetch = dict()
        self._operand_infos = dict()
        self._op_key_to_chunk = defaultdict(list)
        self._terminal_chunk_op_key_to_tileable_key = defaultdict(set)
        self._terminal_tileable_key_to_chunk_op_keys = defaultdict(set)
        self._terminated_tileable_keys = set()
        self._terminated_chunk_keys = set()
        self._terminal_chunk_keys = set()
        self._target_tileable_finished = dict()

        self._assigned_workers = set()
        self._worker_adds = set()
        self._worker_removes = set()
        self._lost_chunks = set()

        self._graph_analyze_pool = None

    def post_create(self):
        super().post_create()
        logger.debug('Actor %s running in process %d', self.uid, os.getpid())

        random.seed(int(time.time()))
        self.set_cluster_info_ref()
        self._assigner_actor_ref = self.get_actor_ref(AssignerActor.gen_uid(self._session_id))
        self._resource_actor_ref = self.get_actor_ref(ResourceActor.default_uid())
        self._session_ref = self.get_actor_ref(SessionActor.gen_uid(self._session_id))

        uid = GraphMetaActor.gen_uid(self._session_id, self._graph_key)
        self._graph_meta_ref = self.ctx.create_actor(
            GraphMetaActor, self._session_id, self._graph_key, self._session_ref, uid=uid)

        self._kv_store_ref = self.ctx.actor_ref(KVStoreActor.default_uid())
        if not self.ctx.has_actor(self._kv_store_ref):
            self._kv_store_ref = None

        self._graph_analyze_pool = self.ctx.threadpool(1)

    def pre_destroy(self):
        super().pre_destroy()
        self.unset_cluster_info_ref()
        self._graph_meta_ref.destroy()

    @contextlib.contextmanager
    def _open_dump_file(self, prefix):  # pragma: no cover
        if options.scheduler.dump_graph_data:
            file_name = f'{prefix}-{self._graph_key}-{int(time.time())}.log'
            yield open(file_name, 'w')
        else:
            try:
                yield None
            except AttributeError:
                return

    @property
    def state(self):
        """
        running state of the graph
        """
        return self._state

    @state.setter
    def state(self, value):
        if value != self._state:
            logger.debug('Graph %s state from %s to %s.', self._graph_key, self._state, value)
        self._state = value
        self._graph_meta_ref.set_state(value, _tell=True, _wait=False)

    @log_unhandled
    def reload_state(self):
        result = self._graph_meta_ref.get_state()
        if result is None:
            return
        state = self._state = result
        return state

    @property
    def final_state(self):
        return self._final_state

    @property
    def context(self):
        if self._context is None:
            self._context = DistributedContext(self.address, self._session_id,
                                               actor_ctx=self.ctx, address=self.address)
        return self._context

    @final_state.setter
    def final_state(self, value):
        self._final_state = value
        self._graph_meta_ref.set_final_state(value, _tell=True, _wait=False)

    def _detect_cancel(self, callback=None):
        if self.reload_state() == GraphState.CANCELLING:
            logger.info('Cancel detected, stopping')
            if callback:
                callback()
            else:
                self._graph_meta_ref.set_graph_end(_tell=True, _wait=False)
                self.state = GraphState.CANCELLED
            raise ExecutionInterrupted

    @log_unhandled
    def execute_graph(self, compose=True):
        """
        Start graph execution
        """

        self._graph_meta_ref.set_graph_start(_tell=True, _wait=False)
        self.state = GraphState.PREPARING
        self._execute_graph(compose=compose)

    def _dump_graph(self):
        with self._open_dump_file('graph') as outf:  # pragma: no cover
            succ_op_keys = dict()
            graph = self._chunk_graph_cache

            outf.write(f'DOT:\n{graph.to_dot()}\n\n')
            outf.write('CONNECTIONS:\n')

            chunks = list(graph)
            for n in graph:
                succ_op_keys[n.key] = succ_op_keys.get(n.key, set()) \
                    | set(succ.op.key for succ in graph.iter_successors(n))
                for c in n.inputs:
                    if isinstance(c.op, Fetch):
                        if c.key not in succ_op_keys:
                            chunks.append(c)
                            succ_op_keys[c.key] = set()
                        succ_op_keys[c.key].add(n.op.key)
            for n in chunks:
                joined_op_keys = ','.join(sorted(succ_op_keys[n.key]))
                outf.write(
                    f'{n.op.key}({n.key})[{type(n.op).__name__}] -> {joined_op_keys}\n'
                )

    def _execute_graph(self, compose=True):
        try:
            self.prepare_graph(compose=compose)
            self._detect_cancel()

            self._dump_graph()

            self.analyze_graph()
            self._detect_cancel()

            if self.state == GraphState.SUCCEEDED:
                self._graph_meta_ref.set_graph_end(_tell=True, _wait=False)
            else:
                self.create_operand_actors()
                self._detect_cancel(self.stop_graph)
        except ExecutionInterrupted:
            pass
        except:  # noqa: E722
            logger.exception('Failed to start graph execution.')
            self._graph_meta_ref.set_exc_info(sys.exc_info(), _tell=True, _wait=False)
            self.stop_graph()
            self.state = GraphState.FAILED
            self._graph_meta_ref.set_graph_end(_tell=True, _wait=False)
            raise

    @log_unhandled
    def stop_graph(self):
        """
        Stop graph execution
        """
        from .operands import OperandActor
        if self.state == GraphState.CANCELLED:
            return
        self.state = GraphState.CANCELLING

        try:
            chunk_graph = self.get_chunk_graph()
        except (KeyError, GraphNotExists, AttributeError):
            self.state = GraphState.CANCELLED
            return

        has_stopping = False
        for chunk in chunk_graph:
            if chunk.op.key not in self._operand_infos:
                continue
            if self._operand_infos[chunk.op.key].get('state') in \
                    (OperandState.READY, OperandState.RUNNING, OperandState.FINISHED):
                # we only need to stop on ready, running and finished operands
                op_uid = OperandActor.gen_uid(self._session_id, chunk.op.key)
                scheduler_addr = self.get_scheduler(op_uid)
                ref = self.ctx.actor_ref(op_uid, address=scheduler_addr)
                has_stopping = True
                ref.stop_operand(_tell=True)
        if not has_stopping:
            self.state = GraphState.CANCELLED
            self._graph_meta_ref.set_graph_end(_tell=True, _wait=False)

    @log_unhandled
    def get_chunk_graph(self):
        try:
            return self._chunk_graph_builder.iterative_chunk_graphs[-1]
        except (IndexError, AttributeError):
            raise GraphNotExists from None

    def _gen_tileable_graph(self):
        if self._tileable_graph_cache is None:
            tileable_graph = deserialize_graph(self._serialized_tileable_graph)
            self._tileable_graph_cache = tileable_graph

            logger.debug('Begin preparing graph %s with %d tileables to chunk graph.',
                         self._graph_key, len(tileable_graph))

        return self._tileable_graph_cache

    def _gen_chunk_graph(self):
        if self._chunk_graph_cache is None:
            self._chunk_graph_cache = DAG()
        return self._chunk_graph_cache

    def _scan_tileable_graph(self):
        tileable_graph = self._tileable_graph_cache
        target_tileable_provided = self._target_tileable_keys is not None
        if self._target_tileable_keys is None:
            self._target_tileable_keys = []
            self._target_tileable_datas = []
        for tn in tileable_graph:
            self._tileable_key_to_opid[tn.key] = tn.op.id
            if tileable_graph.count_successors(tn) == 0:
                # no successors
                if not target_tileable_provided:
                    self._target_tileable_keys.append(tn.key)
                    self._target_tileable_datas.append(tn)
            if target_tileable_provided and tn.key in self._target_tileable_keys:
                self._target_tileable_datas.append(tn)

    @classmethod
    def _merge_chunk_graph(cls, chunk_graph, update_chunk_graph):
        chunk_key_id_to_chunk = dict()
        for c in chunk_graph:
            chunk_key_id_to_chunk[c.key, c.id] = c
        for update_c in update_chunk_graph:
            if isinstance(update_c.op, Fetch):
                continue
            chunk_graph.add_node(update_c)
            for inp in update_c.inputs or []:
                if inp in chunk_graph:
                    chunk_graph.add_edge(inp, update_c)
                elif isinstance(inp.op, Fetch):
                    if (inp.key, inp.id) in chunk_key_id_to_chunk:
                        inp = chunk_key_id_to_chunk[inp.key, inp.id]
                    else:
                        chunk_graph.add_node(inp)
                    chunk_graph.add_edge(inp, update_c)
        return chunk_graph

    @classmethod
    def _prune_chunk_graph(cls, chunk_graph, result_chunk_keys):
        result_op_keys = set()
        for c in chunk_graph:
            if c.key in result_chunk_keys:
                result_op_keys.add(c.op.key)

        reverse_chunk_graph = chunk_graph.build_reversed()
        marked = set()
        for c in reverse_chunk_graph.topological_iter():
            if reverse_chunk_graph.count_predecessors(c) == 0 and \
                    c.op.key in result_op_keys:
                marked.add(c)
            elif any(inp in marked for inp in reverse_chunk_graph.iter_predecessors(c)):
                marked.add(c)
        for n in list(chunk_graph):
            if n not in marked:
                chunk_graph.remove_node(n)

    def _clear_for_iteration(self):
        self._tileable_to_fetch.clear()
        self._op_key_to_chunk.clear()
        self._terminal_chunk_op_key_to_tileable_key.clear()
        self._terminal_tileable_key_to_chunk_op_keys.clear()
        self._target_tileable_finished.clear()
        self._terminated_tileable_keys.clear()
        self._terminal_chunk_keys.clear()
        self._terminated_chunk_keys.clear()

    def _gen_target_info(self):
        chunk_graph_builder = self._chunk_graph_builder
        is_done = chunk_graph_builder.done
        cur_tileable_graph = chunk_graph_builder.prev_tileable_graph
        cur_chunk_graph = self.get_chunk_graph()
        chunk_keys = {c.key for c in cur_chunk_graph}

        if is_done:
            target_tileable_keys = self._target_tileable_keys
        else:
            target_tileable_keys = set(self._target_tileable_keys)
            for interrupted_op in chunk_graph_builder.interrupted_ops:
                for inp in interrupted_op.inputs:
                    if inp.op not in chunk_graph_builder.interrupted_ops:
                        target_tileable_keys.add(inp.key)

        # scan tileable graph
        for tn in cur_tileable_graph:
            if not isinstance(tn.op, Fetch) and tn.key in target_tileable_keys:
                self._target_tileable_finished[tn.key] = set()
                tiled_tn = self._tileable_key_opid_to_tiled[tn.key, tn.op.id][-1]
                for c in tiled_tn.chunks:
                    assert c.key in chunk_keys
                    if isinstance(c.op, Fetch):
                        # if the terminal chunk's op is fetch
                        # just add it to terminated
                        self._terminated_chunk_keys.add(c.key)
                    self._terminal_chunk_keys.add(c.key)
                    self._terminal_chunk_op_key_to_tileable_key[c.op.key].add(tn.key)
                    self._terminal_tileable_key_to_chunk_op_keys[tn.key].add(c.op.key)

        # if not done, scan chunk graph again, for the reason that,
        # make sure some chunks partially generated by tiling of some op
        # are added to terminal chunks correctly
        if not is_done:
            for cn in cur_chunk_graph:
                if cur_chunk_graph.count_successors(cn) == 0 and \
                        not isinstance(cn.op, Fetch):
                    assert cn.key in chunk_keys
                    self._terminal_chunk_keys.add(cn.key)

        logger.debug(f'Terminal chunk keys: {self._terminal_chunk_keys}')

    @staticmethod
    def _get_tileable_graph_builder(graph, **kwargs):
        if options.optimize_tileable_graph:
            return OptimizeIntegratedTileableGraphBuilder(graph, **kwargs)
        else:
            return TileableGraphBuilder(graph, **kwargs)

    @log_unhandled
    @enter_mode(build=True)
    def prepare_graph(self, compose=True):
        """
        Tile and compose tileable graph into chunk graph
        :param compose: if True, do compose after tiling
        """
        tileable_graph = self._gen_tileable_graph()
        chunk_graph = self._gen_chunk_graph()
        self._clear_for_iteration()

        if self._chunk_graph_builder is None:
            # gen target tileable keys if not provided
            self._scan_tileable_graph()

            if options.optimize_tileable_graph:
                target_tileables = [n for n in tileable_graph
                                    if n.key in self._target_tileable_keys]
                tileable_graph = TileableGraph(target_tileables)
                optimized_graph_builder = OptimizeIntegratedTileableGraphBuilder(tileable_graph)
                self._tileable_graph_cache = optimized_graph_builder._build()
                # rescan
                self._target_tileable_datas = list()
                self._scan_tileable_graph()

            def on_tile(raw_tileables, tileds):
                first = tileds[0]
                if any(inp in self._tileable_to_fetch for inp in raw_tileables[0].inputs):
                    new_inputs = []
                    for inp in raw_tileables[0].inputs:
                        if inp in self._tileable_to_fetch:
                            tiled_inp = get_tiled(self._tileable_to_fetch[inp])
                            new_inputs.append(tiled_inp)
                        else:
                            new_inputs.append(get_tiled(inp))
                    first.op.inputs = new_inputs
                if isinstance(first.op, Fetch):
                    raw_first = raw_tileables[0]
                    if (raw_first.key, raw_first.op.id) in self._tileable_key_opid_to_tiled:
                        # fetch op generated by iterative tiling graph builder
                        tiled = self._tileable_key_opid_to_tiled[raw_first.key, raw_first.op.id][-1]
                        return [build_fetch_tileable(tiled)]
                    else:
                        return [self.tile_fetch_tileable(first)]
                else:
                    return self.context.wraps(handler.dispatch)(first.op)

            def on_tile_success(tiled_before, tiled_after):
                if not isinstance(tiled_before.op, Fetch):
                    self._tileable_key_opid_to_tiled[
                        tiled_before.key, tiled_before.op.id].append(tiled_after)
                return tiled_after

            # do not compose here
            self._chunk_graph_builder = IterativeChunkGraphBuilder(
                tileable_graph, on_tile=on_tile, on_tile_success=on_tile_success,
                fuse_enabled=False)

        chunk_graph_builder = self._chunk_graph_builder
        if chunk_graph_builder.prev_tileable_graph is None:
            # first tile
            fetch_tileables = [t for t in tileable_graph if isinstance(t.op, Fetch)]
            # add fetch tileables to make sure that they won't be fused
            chunk_graph_builder._graph._result_tileables = \
                self._target_tileable_datas + fetch_tileables
            cur_chunk_graph = chunk_graph_builder._build()
        else:
            # some TilesFail happens before
            # build tileable graph from failed ops and their inputs
            failed_tileable_set = set(itertools.chain(
                *(op.outputs for op in chunk_graph_builder.interrupted_ops)))
            tileable_graph_builder = self._get_tileable_graph_builder(
                TileableGraph(list(failed_tileable_set)),
                inputs_selector=lambda inps: [inp for inp in inps if inp in failed_tileable_set])
            to_run_tileable_graph = tileable_graph_builder._build()
            to_fetch_tileables = []
            for failed_op in chunk_graph_builder.interrupted_ops:
                for inp in failed_op.inputs:
                    if inp not in failed_tileable_set:
                        if inp not in self._tileable_to_fetch:
                            fetch_inp = build_fetch_tileable(inp).data
                            self._tileable_to_fetch[inp] = fetch_inp
                        else:
                            fetch_inp = self._tileable_to_fetch[inp]
                        to_fetch_tileables.append(fetch_inp)
                        to_run_tileable_graph.add_node(fetch_inp)
                        for o in failed_op.outputs:
                            to_run_tileable_graph.add_edge(fetch_inp, o)
            # add to_fetch_tileables to make sure that fetch chunk would not be fused
            chunk_graph_builder._graph = to_run_tileable_graph
            chunk_graph_builder._graph._result_tileables = \
                self._target_tileable_datas + to_fetch_tileables
            cur_chunk_graph = chunk_graph_builder._build()

        self._gen_target_info()
        if chunk_graph_builder.done:
            self._prune_chunk_graph(cur_chunk_graph, self._terminal_chunk_keys)
        if compose:
            cur_chunk_graph.compose(keys=list(self._terminal_chunk_keys))

        for c in list(cur_chunk_graph):
            if not isinstance(c.op, Fetch):
                self._op_key_to_chunk[c.op.key].append(c)

            # merge current chunk into self.chunk_graph_cache
            if isinstance(c.op, Fetch):
                continue
            chunk_graph.add_node(c)
            self._chunk_key_id_to_chunk[c.key, c.id] = c
            for inp in cur_chunk_graph.iter_predecessors(c):
                if inp in chunk_graph:
                    chunk_graph.add_edge(inp, c)
                elif isinstance(inp.op, Fetch):
                    if (inp.key, inp.id) in self._chunk_key_id_to_chunk:
                        inp = self._chunk_key_id_to_chunk[inp.key, inp.id]
                    else:
                        chunk_graph.add_node(inp)
                    chunk_graph.add_edge(inp, c)

    def _get_worker_slots(self):
        metrics = self._resource_actor_ref.get_workers_meta()
        return dict((ep, int(metrics[ep]['hardware']['cpu_total'])) for ep in metrics)

    def _collect_external_input_metas(self, ext_chunks_to_inputs):
        ext_chunk_keys = reduce(operator.add, ext_chunks_to_inputs.values(), [])
        metas = dict(zip(ext_chunk_keys,
                         self.chunk_meta.batch_get_chunk_meta(self._session_id, ext_chunk_keys)))
        input_chunk_metas = defaultdict(dict)
        for chunk_key, input_chunk_keys in ext_chunks_to_inputs.items():
            chunk_metas = input_chunk_metas[chunk_key]
            for k in input_chunk_keys:
                chunk_metas[k] = metas[k]
        return input_chunk_metas

    def assign_operand_workers(self, op_keys, input_chunk_metas=None, analyzer=None):
        operand_infos = self._operand_infos
        chunk_graph = self.get_chunk_graph()

        initial_chunks = [c for c in chunk_graph
                          if chunk_graph.count_predecessors(c) == 0]
        # TODO refine this to support mixed scenarios here
        if all(c.op.expect_worker is not None for c in initial_chunks):
            assignments = {c.op.key: c.op.expect_worker for c in initial_chunks}
        else:
            if analyzer is None:
                analyzer = GraphAnalyzer(chunk_graph, self._get_worker_slots())
            assignments = analyzer.calc_operand_assignments(op_keys, input_chunk_metas=input_chunk_metas)
        for idx, (k, v) in enumerate(assignments.items()):
            operand_infos[k]['optimize']['placement_order'] = idx
            operand_infos[k]['target_worker'] = v
        return assignments

    def _assign_initial_workers(self, analyzer):
        # collect external inputs for eager mode
        ext_chunks_to_inputs = analyzer.collect_external_input_chunks(initial=True)
        input_chunk_metas = self._collect_external_input_metas(ext_chunks_to_inputs)

        invalid_keys = [k for k, v in input_chunk_metas.items() if v is None]
        if invalid_keys:
            raise KeyError(f'Input metas for data chunks {invalid_keys!r} missing when '
                           f'executing graph {self._graph_key}')

        def _do_assign():
            # do placements
            return self.assign_operand_workers(
                analyzer.get_initial_operand_keys(), input_chunk_metas=input_chunk_metas,
                analyzer=analyzer
            )

        return self._graph_analyze_pool.submit(_do_assign).result()

    @log_unhandled
    def analyze_graph(self, **kwargs):
        operand_infos = self._operand_infos
        chunk_graph = self.get_chunk_graph()

        # remove fetch chunk if exists
        if any(isinstance(c.op, Fetch) for c in chunk_graph):
            chunk_graph = chunk_graph.copy()
            for c in list(chunk_graph):
                if isinstance(c.op, Fetch):
                    chunk_graph.remove_node(c)

        if len(chunk_graph) == 0:
            self.state = GraphState.SUCCEEDED
            return

        for n in chunk_graph:
            k = n.op.key
            succ_size = chunk_graph.count_successors(n)
            if k not in operand_infos:
                operand_infos[k] = dict(optimize=dict(
                    depth=0, demand_depths=(), successor_size=succ_size, descendant_size=0
                ))
            else:
                operand_infos[k]['optimize']['successor_size'] = succ_size

        worker_slots = self._get_worker_slots()
        if not worker_slots:
            raise RuntimeError('No worker attached for execution')

        self._assigned_workers = set(worker_slots)
        analyzer = GraphAnalyzer(chunk_graph, worker_slots)

        for k, v in analyzer.calc_depths().items():
            operand_infos[k]['optimize']['depth'] = v

        for k, v in analyzer.calc_descendant_sizes().items():
            operand_infos[k]['optimize']['descendant_size'] = v

        if kwargs.get('do_placement', True):
            logger.debug('Placing initial chunks for graph %s', self._graph_key)
            self._assign_initial_workers(analyzer)

    @log_unhandled
    def get_executable_operand_dag(self, op_key, input_chunk_keys=None, serialize=True):
        """
        Make an operand into a worker-executable dag
        :param op_key: operand key
        :param input_chunk_keys: actual input chunks, None if use all chunks in input
        :param serialize: whether to return serialized dag
        """
        result_chunks = []
        graph = ChunkGraph(result_chunks)
        input_mapping = dict()
        output_keys = set()

        input_chunk_keys = set(input_chunk_keys) if input_chunk_keys is not None else None
        for c in self._op_key_to_chunk[op_key]:
            inputs = []
            for inp in set(c.op.inputs or ()):
                try:
                    inp_chunk = input_mapping[(inp.key, inp.id)]
                except KeyError:
                    inp_chunk = input_mapping[(inp.key, inp.id)] \
                        = build_fetch_chunk(inp, input_chunk_keys).data
                    graph.add_node(inp_chunk)
                inputs.append(inp_chunk)

            for out in set(c.op.outputs or ()):
                if (out.key, out.id) not in output_keys:
                    output_keys.add((out.key, out.id))
                    graph.add_node(out)
                    for inp in inputs:
                        graph.add_edge(inp, out)
                    result_chunks.append(out)
        if serialize:
            return serialize_graph(graph)
        else:
            return graph

    @staticmethod
    def _collect_operand_io_meta(graph, chunks):
        # collect operand i/o information
        predecessor_keys = set()
        successor_keys = set()
        input_chunk_keys = set()
        shared_input_chunk_keys = set()
        pure_dep_chunk_keys = set()
        chunk_keys = set()
        virtual_chunk_keys = set()
        reducer_indexes = dict()
        predecessors_to_successors = dict()

        for c in chunks:
            # handling predecessor args
            for pn in graph.iter_predecessors(c):
                if isinstance(pn.op, VirtualOperand):
                    virtual_chunk_keys.add(pn.key)
                if not isinstance(pn.op, Fetch):
                    predecessor_keys.add(pn.op.key)
                input_chunk_keys.add(pn.key)
                if graph.count_successors(pn) > 1:
                    shared_input_chunk_keys.add(pn.key)

            for inp, is_dep in zip(c.op.inputs or (), c.op.pure_depends):
                if is_dep and inp.key in input_chunk_keys:
                    pure_dep_chunk_keys.add(inp.key)

            # handling successor args
            for sn in graph.iter_successors(c):
                successor_keys.add(sn.op.key)
            if isinstance(c.op, ShuffleProxy):
                for sn in graph.iter_successors(c):
                    reducer_indexes[sn.op.key] = get_chunk_reducer_index(sn)
            if isinstance(c.op, SuccessorsExclusive):
                for sn in graph.iter_successors(c):
                    predecessors_to_successors[sn.inputs[0].op.key] = sn.op.key

            chunk_keys.update(co.key for co in c.op.outputs)

        io_meta = dict(
            predecessors=list(predecessor_keys),
            successors=list(successor_keys),
            input_chunks=list(input_chunk_keys),
            pure_dep_chunk_keys=list(pure_dep_chunk_keys),
            shared_input_chunks=list(shared_input_chunk_keys),
            chunks=list(chunk_keys),
        )
        if virtual_chunk_keys:
            io_meta['virtual_chunk_keys'] = list(virtual_chunk_keys)
        if reducer_indexes:
            io_meta['reducer_indexes'] = [reducer_indexes.get(k) for k in io_meta['successors']]
        if predecessors_to_successors:
            io_meta['predecessors_to_successors'] = predecessors_to_successors
        return io_meta

    @log_unhandled
    def create_operand_actors(self, _clean_info=True, _start=True):
        """
        Create operand actors for all operands
        """
        logger.debug('Creating operand actors for graph %s', self._graph_key)

        chunk_graph = self.get_chunk_graph()
        operand_infos = self._operand_infos

        session_id = self._session_id
        op_refs = dict()
        meta_op_infos = dict()
        initial_keys = []
        to_allocate_op_keys = set()

        total_ops = len(self._op_key_to_chunk)
        processed_ops = 0
        last_progress = 0
        for op_key in self._op_key_to_chunk:
            chunks = self._op_key_to_chunk[op_key]
            op = chunks[0].op

            if isinstance(op, VirtualOperand):
                operand_infos[op_key]['virtual'] = True

            op_name = type(op).__name__ if op.stage is None \
                else f'{type(op).__name__}:{op.stage.name}'
            op_info = operand_infos[op_key]
            meta_op_info = meta_op_infos[op_key] = dict()

            io_meta = self._collect_operand_io_meta(chunk_graph, chunks)
            op_info['op_name'] = meta_op_info['op_name'] = op_name
            op_info['io_meta'] = io_meta
            op_info['executable_dag'] = self._graph_analyze_pool.submit(
                self.get_executable_operand_dag, op_key).result()
            # todo change this when other calc devices supported
            op_info['calc_device'] = 'cuda' if op.gpu else 'cpu'

            if io_meta['predecessors']:
                state = OperandState.UNSCHEDULED
            else:
                initial_keys.append(op_key)
                state = OperandState.READY
                op_info['is_initial'] = True
            op_info['retryable'] = op.retryable
            op_info['retries'] = 0
            meta_op_info['state'] = op_info['state'] = state
            meta_op_info['worker'] = op_info.get('worker')

            if any(c.key in self._terminal_chunk_keys for c in op.outputs):
                op_info['is_terminal'] = True

            op_cls = get_operand_actor_class(type(op))
            op_uid = op_cls.gen_uid(session_id, op_key)
            scheduler_addr = self.get_scheduler(op_uid)
            kw = {}
            # for the **real** initial chunks, we do batch submitting
            if options.scheduler.batch_enqueue_initials \
                    and chunk_graph.count_predecessors(chunks[0]) == 0:
                kw['allocated'] = True
                to_allocate_op_keys.add(op_key)

            op_refs[op_key] = self.ctx.create_actor(
                op_cls, session_id, self._graph_key, op_key, op_info.copy(),
                with_kvstore=self._kv_store_ref is not None,
                schedulers=self.get_schedulers(),
                uid=op_uid, address=scheduler_addr, wait=False, **kw
            )
            if _clean_info:
                op_info.pop('executable_dag', None)
                del op_info['io_meta']

            processed_ops += 1
            progress = processed_ops * 1.0 / total_ops
            if int(progress * 20) > last_progress:
                last_progress = int(progress * 20)
                logger.info('Operand actor creation progress: %d / %d', processed_ops, total_ops)

        self.state = GraphState.RUNNING
        self._graph_meta_ref.update_op_infos(meta_op_infos, _tell=True, _wait=False)

        if _start:
            existing_keys = []
            created_keys = []
            for op_key, future in op_refs.items():
                try:
                    op_refs[op_key] = future.result()
                    created_keys.append(op_key)
                except ActorAlreadyExist:
                    existing_keys.append(op_key)

            # start operands when all operands are created
            for op_key in created_keys:
                op_refs[op_key].start_operand(_tell=True, _wait=False)

            append_futures = []
            for op_key in existing_keys:
                chunks = self._op_key_to_chunk[op_key]
                op = chunks[0].op
                op_info = operand_infos[op_key]
                op_info['io_meta'] = self._collect_operand_io_meta(chunk_graph, chunks)

                op_cls = get_operand_actor_class(type(op))
                op_uid = op_cls.gen_uid(session_id, op_key)
                scheduler_addr = self.get_scheduler(op_uid)
                op_ref = op_refs[op_key] = self.ctx.actor_ref(op_uid, address=scheduler_addr)
                append_futures.append(op_ref.append_graph(self._graph_key, op_info.copy(), _wait=False))
                if _clean_info:
                    op_info.pop('executable_dag', None)
                    del op_info['io_meta']
            [future.result() for future in append_futures]

            res_applications = [(op_key, op_info) for op_key, op_info in operand_infos.items()
                                if op_key in to_allocate_op_keys]
            self._assigner_actor_ref.apply_for_multiple_resources(
                session_id, res_applications, _tell=True)

    @log_unhandled
    def add_finished_terminal(self, op_key, final_state=None, exc=None):
        """
        Add a terminal operand to finished set. Calling this method
        will change graph state if all terminals are in finished states.
        :param op_key: operand key
        :param final_state: state of the operand
        """
        if self._state not in (GraphState.RUNNING, GraphState.CANCELLING):
            return
        if exc is not None:
            self._graph_meta_ref.set_exc_info(exc, _tell=True, _wait=False)

        tileable_keys = self._terminal_chunk_op_key_to_tileable_key[op_key]
        is_failed = final_state in (GraphState.CANCELLED, GraphState.FAILED)
        terminal_tileable_count = len(self._terminal_tileable_key_to_chunk_op_keys)
        try:
            for tileable_key in tileable_keys:
                self._target_tileable_finished[tileable_key].add(op_key)
                if final_state == GraphState.FAILED:
                    if self.final_state != GraphState.CANCELLED:
                        self.final_state = GraphState.FAILED
                elif final_state == GraphState.CANCELLED:
                    self.final_state = final_state

                if self._target_tileable_finished[tileable_key] == \
                        self._terminal_tileable_key_to_chunk_op_keys[tileable_key]:
                    self._terminated_tileable_keys.add(tileable_key)
                    self._all_terminated_tileable_keys.add(tileable_key)
                    if not is_failed and len(self._terminated_tileable_keys) == terminal_tileable_count:
                        # update shape if tileable or its chunks have unknown shape
                        self._update_tileable_and_its_chunk_shapes()
        except:
            for tileable_key in tileable_keys:
                self._target_tileable_finished[tileable_key].remove(op_key)
            raise

        terminated_chunks = self._op_key_to_chunk[op_key]
        self._terminated_chunk_keys.update([c.key for c in terminated_chunks
                                            if c.key in self._terminal_chunk_keys])
        if self._terminated_chunk_keys == self._terminal_chunk_keys:
            if self._chunk_graph_builder.done or is_failed:
                if self._chunk_graph_builder.prev_tileable_graph is not None:
                    # if failed before, clear intermediate data
                    to_free_tileable_keys = \
                        self._all_terminated_tileable_keys - set(self._target_tileable_keys)
                    skip_chunk_keys = set()
                    for target_tileable_data in self._target_tileable_datas:
                        tiled_target_tileable_data = \
                            self._tileable_key_opid_to_tiled[target_tileable_data.key,
                                                             target_tileable_data.op.id][-1]
                        skip_chunk_keys.update([c.key for c in tiled_target_tileable_data.chunks])
                    [self.free_tileable_data(k, skip_chunk_keys=skip_chunk_keys)
                     for k in to_free_tileable_keys]
                self.state = self.final_state if self.final_state is not None else GraphState.SUCCEEDED
                self._graph_meta_ref.set_graph_end(_tell=True)
            else:
                self._execute_graph(compose=self._chunk_graph_builder.fused_enabled)

    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)

    @log_unhandled
    def remove_finished_terminal(self, op_key):
        """
        Remove a terminal operand from finished set as the data is lost.
        :param op_key: operand key
        """
        tileable_keys = self._terminal_chunk_op_key_to_tileable_key[op_key]
        for tileable_key in tileable_keys:
            self._target_tileable_finished[tileable_key].difference_update([op_key])
            self._terminated_tileable_keys.difference_update([tileable_key])

    def dump_unfinished_terminals(self):  # pragma: no cover
        """
        Dump unfinished terminal chunks into logger, only for debug purposes.
        """
        unfinished_dict = dict()
        for tileable_key, chunk_ops in self._terminal_tileable_key_to_chunk_op_keys.items():
            executed_ops = self._target_tileable_finished.get(tileable_key, ())
            unfinished = sorted(set(chunk_ops) - set(executed_ops))
            if unfinished:
                unfinished_dict[tileable_key] = unfinished
        logger.debug('Unfinished terminal chunks: %r', unfinished_dict)

    def get_state(self):
        return self.state

    def get_operand_states(self, op_keys):
        return [self._operand_infos[k]['state'] for k in op_keys if k in self._operand_infos]

    def set_operand_state(self, op_key, state):
        if op_key not in self._operand_infos and \
                self._chunk_graph_builder.iterative_chunk_graphs and \
                state == OperandState.FREED:
            # if iterative tiling is entered,
            # the `_operand_infos` will be a completely new one,
            # in this case, we don't actually care about if the op is freed
            return
        if op_key not in self._operand_infos and self.state in GraphState.TERMINATED_STATES:
            # if operand has been cleared in iterative tiling and execute again in another
            # graph, just ignore it.
            return
        op_info = self._operand_infos[op_key]
        op_info['state'] = state
        self._graph_meta_ref.update_op_state(op_key, op_info['op_name'], state,
                                             _tell=True, _wait=False)
        op_info.pop('failover_state', None)

    def get_operand_target_worker(self, op_key):
        return self._operand_infos[op_key]['target_worker']

    def get_operand_info(self, state=None):
        if not state:
            return self._operand_infos
        else:
            return dict((k, info) for k, info in self._operand_infos.items()
                        if info.get('state') == state)

    @log_unhandled
    def set_operand_worker(self, op_key, worker):
        if op_key not in self._operand_infos and self.state in GraphState.TERMINATED_STATES:
            # if operand has been cleared in iterative tiling and execute again in another
            # graph, just ignore it.
            return
        op_info = self._operand_infos[op_key]
        if worker:
            op_info['worker'] = worker
        else:
            op_info.pop('worker', None)
        self._graph_meta_ref.update_op_worker(op_key, op_info['op_name'], worker,
                                              _tell=True, _wait=False)

    @lru_cache(1000)
    def _get_operand_ref(self, key):
        from .operands import OperandActor
        op_uid = OperandActor.gen_uid(self._session_id, key)
        scheduler_addr = self.get_scheduler(op_uid)
        return self.ctx.actor_ref(op_uid, address=scheduler_addr)

    def _get_tileable_by_key(self, key):
        tid = self._tileable_key_to_opid[key]
        return self._tileable_key_opid_to_tiled[(key, tid)][-1]

    @log_unhandled
    def free_tileable_data(self, tileable_key, skip_chunk_keys=None, wait=False):
        tileable = self._get_tileable_by_key(tileable_key)
        futures = []
        for chunk in tileable.chunks:
            if skip_chunk_keys is not None and chunk.key in skip_chunk_keys:
                continue
            if isinstance(chunk.op, Fetch):
                # skip fetch op
                continue
            futures.append(self._get_operand_ref(chunk.op.key).free_data(
                check=False, _tell=not wait, _wait=False))
        logger.debug('Free tileable data: %s, chunk keys: %r',
                     tileable_key, [c.key for c in tileable.chunks])
        [f.result() for f in futures]

    def get_tileable_metas(self, tileable_keys, filter_fields=None):
        """
        Get tileable meta including nsplits, chunk keys and chunk indexes.
        :param tileable_keys: tileable_keys.
        :param filter_fields: filter the fields('nsplits', 'chunk_keys', 'chunk_indexes') in meta.
        :return: metas list.
        """
        meta_names = ['nsplits', 'chunk_keys', 'chunk_indexes', 'extra_meta']
        filter_fields = filter_fields or meta_names
        for field in filter_fields:
            if field not in ['nsplits', 'chunk_keys', 'chunk_indexes', 'extra_meta']:
                raise ValueError(f'Field {field} is invalid')

        metas = []
        for tileable_key in tileable_keys:
            meta = dict()
            tileable = self._get_tileable_by_key(tileable_key)
            chunk_keys, chunk_indexes = tuple(zip(*[(c.key, c.index) for c in tileable.chunks]))
            meta['chunk_keys'] = chunk_keys
            meta['chunk_indexes'] = chunk_indexes

            if hasattr(tileable, 'shape') and np.nan in tileable.shape:
                chunk_shapes = self.chunk_meta.batch_get_chunk_shape(self._session_id,  chunk_keys)
                nsplits = calc_nsplits(OrderedDict(zip(chunk_indexes, chunk_shapes)))
            else:
                nsplits = tileable.nsplits
            meta['nsplits'] = nsplits
            extra_meta = tileable.params
            extra_meta.pop('shape', None)
            extra_meta = dict((k, v.to_pandas() if hasattr(v, 'to_pandas') else v)
                              for k, v in extra_meta.items())
            meta['extra_meta'] = extra_meta
            metas.append([meta[k] for k in filter_fields])
        return metas

    def get_tileable_chunk_metas(self, tileable_keys, filter_fields=None):
        metas = []
        for tileable_key in tileable_keys:
            tileable = self._get_tileable_by_key(tileable_key)
            chunk_keys = [c.key for c in tileable.chunks]
            meta = self.chunk_meta.batch_get_chunk_meta(
                self._session_id, chunk_keys, filter_fields=filter_fields)
            metas.append({k: m for k, m in zip(chunk_keys, meta)})

        return metas

    def build_fetch_graph(self, tileable_key):
        """
        Convert single tileable node to tiled fetch tileable node and
        put into a graph which only contains one tileable node
        :param tileable_key: the key of tileable node
        """
        tileable = self._get_tileable_by_key(tileable_key)
        results = []
        graph = TileableGraph(results)

        new_tileable = build_fetch_tileable(tileable).data
        graph.add_node(new_tileable)
        results.append(new_tileable)
        return serialize_graph(graph)

    def tile_fetch_tileable(self, tileable):
        """
        Find the owner of the input tileable node and ask for tiling
        """
        tileable_key = tileable.key
        graph_ref = self.ctx.actor_ref(self._session_ref.get_graph_ref_by_tileable_key(tileable_key))
        fetch_graph = deserialize_graph(graph_ref.build_fetch_graph(tileable_key))
        return list(fetch_graph)[0]

    @log_unhandled
    @enter_mode(build=True)
    def fetch_tileable_result(self, tileable_key, _check=True):
        # this function is for test usage
        from ..executor import Executor
        from ..worker.transfer import ResultSenderActor

        tileable = self._get_tileable_by_key(tileable_key)
        if _check and tileable_key not in self._terminated_tileable_keys:
            return None

        ctx = dict()
        for chunk in tileable.chunks:
            endpoints = self.chunk_meta.get_workers(self._session_id, chunk.key)
            if endpoints is None:
                raise KeyError(f'cannot fetch meta of chunk {chunk}')
            sender_ref = self.ctx.actor_ref(ResultSenderActor.default_uid(), address=endpoints[-1])
            ctx[chunk.key] = dataserializer.loads(sender_ref.fetch_data(self._session_id, chunk.key))

        if len(tileable.chunks) == 1:
            return dataserializer.dumps(ctx[tileable.chunks[0].key])

        fetch_tileable = build_fetch_tileable(tileable)
        concat_chunk = fetch_tileable.op.concat_tileable_chunks(fetch_tileable).chunks[0].data
        chunk_graph = DAG()
        chunk_graph.add_node(concat_chunk)
        for fetch_chunk in fetch_tileable.chunks:
            chunk_graph.add_node(fetch_chunk.data)
            chunk_graph.add_edge(fetch_chunk.data, concat_chunk)
        executor = Executor(storage=ctx)
        concat_result = executor.execute_graph(chunk_graph, keys=[concat_chunk.key])[0]
        return dataserializer.dumps(concat_result)

    @log_unhandled
    def add_fetch_tileable(self, tileable_key, tileable_id, shape, dtype, chunk_size, chunk_keys):
        from ..tensor.utils import create_fetch_tensor
        tensor = create_fetch_tensor(chunk_size, shape, dtype,
                                     tileable_key, tileable_id, chunk_keys)
        self._tileable_key_to_opid[tileable_key] = tensor.op.id
        self._tileable_key_opid_to_tiled[(tileable_key, tensor.op.id)].append(tensor)

    @log_unhandled
    def check_operand_can_be_freed(self, succ_op_keys):
        """
        Check if the data of an operand can be freed.

        :param succ_op_keys: keys of successor operands
        :return: True if can be freed, False if cannot. None when the result
                 is not determinant and we need to test later.
        """
        operand_infos = self._operand_infos
        for k in succ_op_keys:
            if k not in operand_infos:
                continue
            op_info = operand_infos[k]
            op_state = op_info.get('state')
            if op_state not in OperandState.SUCCESSFUL_STATES:
                return False
            failover_state = op_info.get('failover_state')
            if failover_state and failover_state not in OperandState.SUCCESSFUL_STATES:
                return False
        # if can be freed but blocked by an ongoing fail-over step,
        # we try later.
        if self._operand_free_paused:
            return None
        return True

    @log_unhandled
    def handle_worker_change(self, adds, removes, lost_chunks, handle_later=True):
        """
        Calculate and propose changes of operand states given changes
        in workers and lost chunks.

        :param adds: endpoints of workers newly added to the cluster
        :param removes: endpoints of workers removed to the cluster
        :param lost_chunks: keys of lost chunks
        :param handle_later: run the function later, only used in this actor
        """
        if self._state in GraphState.TERMINATED_STATES:
            return

        if handle_later:
            # Run the fail-over process later.
            # This is the default behavior as we need to make sure that
            # all crucial state changes are received by GraphActor.
            # During the delay, no operands are allowed to be freed.
            self._operand_free_paused = True
            self._worker_adds.update(adds)
            self._worker_removes.update(removes)
            self._lost_chunks.update(lost_chunks)
            self.ref().handle_worker_change(adds, removes, lost_chunks,
                                            handle_later=False, _delay=0.5, _tell=True)
            return
        else:
            self._operand_free_paused = False

        adds = self._worker_adds
        self._worker_adds = set()
        removes = self._worker_removes
        self._worker_removes = set()
        lost_chunks = self._lost_chunks
        self._lost_chunks = set()
        if not adds and not removes:
            return

        if all(ep in self._assigned_workers for ep in adds) \
                and not any(ep in self._assigned_workers for ep in removes):
            return

        if not options.scheduler.enable_failover:
            if removes:
                logger.exception('Failed as worker dead and fail-over not enabled.')
                self._graph_meta_ref.set_exc_info(build_exc_info(WorkerDead), _tell=True, _wait=False)
                self.stop_graph()
                self.state = GraphState.FAILED
                self._graph_meta_ref.set_graph_end(_tell=True, _wait=False)
            else:  # pragma: no cover
                logger.warning('New workers added. Existing operands will not be rescheduled.')
            return

        worker_slots = self._get_worker_slots()
        self._assigned_workers = set(worker_slots)
        removes_set = set(removes)

        # collect operand states
        operand_infos = self._operand_infos
        fixed_assigns = dict()
        graph_states = dict()
        for key, op_info in operand_infos.items():
            op_worker = op_info.get('worker')
            if op_worker is None:
                continue

            op_state = graph_states[key] = op_info['state']

            # RUNNING nodes on dead workers should be moved to READY first
            if op_state == OperandState.RUNNING and op_worker in removes_set:
                graph_states[key] = OperandState.READY

            if op_worker in worker_slots:
                fixed_assigns[key] = op_info['worker']

        graph = self.get_chunk_graph()
        new_states = dict()
        ordered_states = OrderedDict(
            sorted(((k, v) for k, v in graph_states.items()),
                   key=lambda d: operand_infos[d[0]]['optimize'].get('placement_order', 0))
        )
        analyzer = GraphAnalyzer(graph, worker_slots, fixed_assigns, ordered_states, lost_chunks)
        if removes or lost_chunks:
            new_states = analyzer.analyze_state_changes()
            logger.debug('%d chunks lost. %d operands changed state.', len(lost_chunks),
                         len(new_states))

        logger.debug('Start reallocating initial operands')
        new_targets = dict(self._assign_initial_workers(analyzer))

        futures = []
        # make sure that all readies and runnings are included to be checked
        for key, op_info in operand_infos.items():
            if key in new_states:
                continue
            state = op_info['state']
            if state == OperandState.RUNNING and \
                    operand_infos[key]['worker'] not in removes_set:
                continue
            if state in (OperandState.READY, OperandState.RUNNING):
                new_states[key] = state

        for key, state in new_states.items():
            new_target = new_targets.get(key)

            op_info = operand_infos[key]
            from_state = op_info['state']
            # record the target state in special info key
            # in case of concurrency issues
            op_info['failover_state'] = state

            op_ref = self._get_operand_ref(key)
            # states may easily slip into the next state when we are
            # calculating fail-over states. Hence we need to include them
            # into source states.
            if from_state == OperandState.READY:
                from_states = [from_state, OperandState.RUNNING]
            elif from_state == OperandState.RUNNING:
                from_states = [from_state, OperandState.FINISHED]
            elif from_state == OperandState.FINISHED:
                from_states = [from_state, OperandState.FREED]
            else:
                from_states = [from_state]
            futures.append(op_ref.move_failover_state(
                from_states, state, new_target, removes, _tell=True, _wait=False))
        [f.result() for f in futures]

        self._dump_failover_info(adds, removes, lost_chunks, new_states)

    def _dump_failover_info(self, adds, removes, lost_chunks, new_states):  # pragma: no cover
        if not options.scheduler.dump_graph_data:
            return
        with self._open_dump_file('failover-record') as outf:
            outf.write('ADDED WORKERS:\n')
            for c in adds:
                outf.write(c + '\n')
            outf.write('REMOVED WORKERS:\n')
            for c in removes:
                outf.write(c + '\n')
            outf.write('\n\nLOST CHUNKS:\n')
            for c in lost_chunks:
                outf.write(repr(c) + '\n')
            outf.write('\n\nOPERAND SNAPSHOT:\n')
            for key, op_info in self._operand_infos.items():
                outf.write(f'Chunk: {key} Worker: {op_info.get("worker")} '
                           f'State: {op_info["state"].value}\n')
            outf.write('\n\nSTATE TRANSITIONS:\n')
            for key, state in new_states.items():
                outf.write(f'{key} -> {state.name}\n')
