# 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 functools
import logging
import random
import time
from collections import defaultdict

from .. import promise
from ..config import options
from ..core.operand import Fetch, FetchShuffle
from ..errors import PinDataKeyFailed, WorkerProcessStopped, WorkerDead, \
    ExecutionInterrupted, DependencyMissing
from ..executor import Executor
from ..utils import BlacklistSet, deserialize_graph, log_unhandled, build_exc_info, \
    calc_data_size, get_chunk_reducer_index, calc_object_overhead
from .storage import DataStorageDevice
from .transfer import ReceiverManagerActor
from .utils import WorkerActor, ExpiringCache, ExecutionState, concat_operand_keys, \
    build_quota_key, change_quota_key_owner

logger = logging.getLogger(__name__)


class GraphExecutionRecord(object):
    """
    Execution records of the graph
    """
    __slots__ = ('graph', 'graph_serialized', '_state', 'op_string', 'data_targets',
                 'chunk_targets', 'io_meta', 'data_metas', 'shared_input_chunks',
                 'state_time', 'mem_request', 'pinned_keys', 'mem_overhead_keys',
                 'est_finish_time', 'calc_actor_uid', 'send_addresses', 'retry_delay',
                 'retry_pending', 'finish_callbacks', 'stop_requested', 'calc_device',
                 'preferred_data_device', 'resource_released', 'pure_dep_chunk_keys')

    def __init__(self, graph_serialized, state, chunk_targets=None, data_targets=None,
                 io_meta=None, data_metas=None, mem_request=None, shared_input_chunks=None,
                 pinned_keys=None, mem_overhead_keys=None, est_finish_time=None,
                 calc_actor_uid=None, send_addresses=None, retry_delay=None,
                 finish_callbacks=None, stop_requested=False, calc_device=None,
                 preferred_data_device=None, resource_released=False,
                 pure_dep_chunk_keys=None):

        self.graph_serialized = graph_serialized
        graph = self.graph = deserialize_graph(graph_serialized)

        self._state = state
        self.state_time = time.time()
        self.data_targets = data_targets or []
        self.chunk_targets = chunk_targets or []
        self.io_meta = io_meta or dict()
        self.data_metas = data_metas or dict()
        self.shared_input_chunks = shared_input_chunks or set()
        self.mem_request = mem_request or dict()
        self.pinned_keys = set(pinned_keys or [])
        self.mem_overhead_keys = set(mem_overhead_keys or [])
        self.est_finish_time = est_finish_time or time.time()
        self.calc_actor_uid = calc_actor_uid
        self.send_addresses = send_addresses
        self.retry_delay = retry_delay or 0
        self.retry_pending = False
        self.finish_callbacks = finish_callbacks or []
        self.stop_requested = stop_requested or False
        self.calc_device = calc_device
        self.preferred_data_device = preferred_data_device
        self.resource_released = resource_released
        self.pure_dep_chunk_keys = pure_dep_chunk_keys or set()

        _, self.op_string = concat_operand_keys(graph)

    @property
    def state(self):
        return self._state

    @state.setter
    def state(self, value):
        self._state = value
        self.state_time = time.time()


class GraphResultRecord(object):
    """
    Execution result of a graph
    """
    __slots__ = 'data_sizes', 'data_shapes', 'exc', 'succeeded'

    def __init__(self, *args, **kwargs):
        succeeded = self.succeeded = kwargs.pop('succeeded', True)
        if succeeded:
            self.data_sizes = args[0]
            self.data_shapes = args[1]
        else:
            self.exc = args

    def build_args(self):
        if self.succeeded:
            return (self.data_sizes, self.data_shapes), {}
        else:
            return self.exc, dict(_accept=False)


class ExecutionActor(WorkerActor):
    """
    Actor for execution control
    """
    _last_dump_time = time.time()

    def __init__(self):
        super().__init__()
        self._dispatch_ref = None
        self._mem_quota_ref = None
        self._status_ref = None
        self._daemon_ref = None
        self._receiver_manager_ref = None

        self._resource_ref = None

        self._graph_records = dict()  # type: dict[tuple, GraphExecutionRecord]
        self._result_cache = ExpiringCache()  # type: dict[tuple, GraphResultRecord]

        self._peer_blacklist = BlacklistSet(options.worker.peer_blacklist_time)

    def post_create(self):
        from .daemon import WorkerDaemonActor
        from .dispatcher import DispatchActor
        from .quota import MemQuotaActor
        from .status import StatusActor

        super().post_create()

        self._dispatch_ref = self.promise_ref(DispatchActor.default_uid())
        self._mem_quota_ref = self.promise_ref(MemQuotaActor.default_uid())

        self._daemon_ref = self.ctx.actor_ref(WorkerDaemonActor.default_uid())
        if not self.ctx.has_actor(self._daemon_ref):
            self._daemon_ref = None
        else:
            self.register_actors_down_handler()

        self._status_ref = self.ctx.actor_ref(StatusActor.default_uid())
        if not self.ctx.has_actor(self._status_ref):
            self._status_ref = None

        self._receiver_manager_ref = self.ctx.actor_ref(ReceiverManagerActor.default_uid())
        if not self.ctx.has_actor(self._receiver_manager_ref):
            self._receiver_manager_ref = None
        else:
            self._receiver_manager_ref = self.promise_ref(self._receiver_manager_ref)

        from ..scheduler import ResourceActor
        self._resource_ref = self.get_actor_ref(ResourceActor.default_uid())

        self.periodical_dump()

    def periodical_dump(self):
        """
        Periodically dump debug information
        """
        if logger.getEffectiveLevel() > logging.DEBUG:
            return
        cls = type(self)
        if cls._last_dump_time < time.time() - 10:
            cls._last_dump_time = time.time()
            if self._graph_records:
                self._dump_execution_states()
        self.ref().periodical_dump(_tell=True, _delay=10)

    def _pin_shared_data_keys(self, session_id, graph_key, data_keys):
        if not data_keys:
            return []
        try:
            graph_record = self._graph_records[(session_id, graph_key)]
        except KeyError:
            return []
        pinned_keys = self.storage_client.pin_data_keys(
            session_id, data_keys, graph_key, [graph_record.preferred_data_device])
        graph_record.pinned_keys.update(pinned_keys)
        return pinned_keys

    def _estimate_calc_memory(self, session_id, graph_key):
        graph_record = self._graph_records[(session_id, graph_key)]
        data_metas = graph_record.data_metas
        size_ctx = dict((k, (v.chunk_size, v.chunk_size)) for k, v in data_metas.items())

        # update shapes
        for n in graph_record.graph:
            if isinstance(n.op, Fetch):
                try:
                    meta = data_metas[n.key]
                    if hasattr(n, '_shape') and meta.chunk_shape is not None:
                        n._shape = meta.chunk_shape
                except KeyError:
                    pass
            elif isinstance(n.op, FetchShuffle):
                reducer_index = get_chunk_reducer_index(graph_record.graph.successors(n)[0])
                shapes = [data_metas.get((k, reducer_index)).chunk_shape for k in n.op.source_keys]
                n.extra_params['_shapes'] = \
                    dict(((k, reducer_index), v) for k, v in zip(n.op.source_keys, shapes))

        executor = Executor(storage=size_ctx, sync_provider_type=Executor.SyncProviderType.MOCK)
        res = executor.execute_graph(graph_record.graph, graph_record.chunk_targets, mock=True)

        targets = graph_record.chunk_targets
        target_sizes = dict(zip(targets, res))

        total_mem = sum(target_sizes[key][1] for key in targets)
        if total_mem:
            for key in targets:
                r = target_sizes[key]
                target_sizes[key] = (r[0], max(r[1], r[1] * executor.mock_max_memory // total_mem))
        return target_sizes

    def _prepare_quota_request(self, session_id, graph_key):
        """
        Calculate quota request for an execution graph
        :param session_id: session id
        :param graph_key: key of the execution graph
        :return: allocation dict
        """
        try:
            graph_record = self._graph_records[(session_id, graph_key)]
        except KeyError:
            return None

        graph = graph_record.graph
        storage_client = self.storage_client
        input_data_sizes = dict((k, v.chunk_size) for k, v in graph_record.data_metas.items())
        alloc_mem_batch = dict()
        alloc_cache_batch = dict()
        input_chunk_keys = dict()

        if self._status_ref:
            self.estimate_graph_finish_time(session_id, graph_key)

        if graph_record.preferred_data_device == DataStorageDevice.SHARED_MEMORY or \
                graph_record.preferred_data_device == DataStorageDevice.VINEYARD:  # pragma: no cover
            memory_estimations = self._estimate_calc_memory(session_id, graph_key)
        else:
            memory_estimations = dict()

        graph_record.mem_overhead_keys = set()
        # collect potential allocation sizes
        for chunk in graph:
            op = chunk.op
            overhead_keys_and_shapes = []

            if isinstance(op, Fetch):
                if chunk.key in chunk.key in graph_record.pure_dep_chunk_keys:
                    continue
                # use actual size as potential allocation size
                input_chunk_keys[chunk.key] = input_data_sizes.get(chunk.key) or calc_data_size(chunk)
                overhead_keys_and_shapes = [(chunk.key, getattr(chunk, 'shape', None))]
            elif isinstance(op, FetchShuffle):
                reducer_index = get_chunk_reducer_index(graph.successors(chunk)[0])
                overhead_keys_and_shapes = chunk.extra_params.get('_shapes', dict()).items()
                for k in op.source_keys:
                    part_key = (k, reducer_index)
                    try:
                        input_chunk_keys[part_key] = input_data_sizes[part_key]
                    except KeyError:
                        pass
            elif chunk.key in graph_record.chunk_targets:
                # use estimated size as potential allocation size
                estimation_data = memory_estimations.get(chunk.key)
                if not estimation_data:
                    continue
                quota_key = build_quota_key(session_id, chunk.key, owner=graph_key)
                cache_batch, alloc_mem_batch[quota_key] = estimation_data
                if not isinstance(chunk.key, tuple):
                    alloc_cache_batch[chunk.key] = cache_batch

            for key, shape in overhead_keys_and_shapes:
                overhead = calc_object_overhead(chunk, shape)
                if overhead:
                    graph_record.mem_overhead_keys.add(key)
                    quota_key = build_quota_key(session_id, key, owner=graph_key)
                    alloc_mem_batch[quota_key] = overhead

        keys_to_pin = list(input_chunk_keys.keys())
        graph_record.pinned_keys = set()
        self._pin_shared_data_keys(session_id, graph_key, keys_to_pin)

        for k, v in input_chunk_keys.items():
            quota_key = build_quota_key(session_id, k, owner=graph_key)
            if quota_key not in alloc_mem_batch:
                if k in graph_record.pinned_keys or k in graph_record.shared_input_chunks:
                    continue
            alloc_mem_batch[quota_key] = alloc_mem_batch.get(quota_key, 0) + v

        if alloc_cache_batch:
            storage_client.spill_size(sum(alloc_cache_batch.values()), [graph_record.preferred_data_device])

        graph_record.mem_request = alloc_mem_batch or dict()
        return alloc_mem_batch

    @log_unhandled
    def _fetch_remote_data(self, session_id, graph_key, chunk_key, remote_addr, *_, **kwargs):
        """
        Asynchronously send data receiving command to a remote address
        :param session_id: session id
        :param graph_key: graph key
        :param chunk_key: chunk key
        :param remote_addr: remote server containing provided chunk key
        :return: promise object
        """
        from .dispatcher import DispatchActor

        if remote_addr in self._peer_blacklist:
            raise DependencyMissing((session_id, graph_key))

        remote_disp_ref = self.promise_ref(uid=DispatchActor.default_uid(),
                                           address=remote_addr)
        ensure_cached = kwargs.pop('ensure_cached', True)
        timeout = options.worker.prepare_data_timeout
        storage_client = self.storage_client
        graph_record = self._graph_records[(session_id, graph_key)]

        @log_unhandled
        def _finish_fetch(*_):
            locations = set(locs[1] for locs in storage_client.get_data_locations(session_id, [chunk_key])[0])
            if DataStorageDevice.PROC_MEMORY not in locations:
                self._pin_shared_data_keys(session_id, graph_key, [chunk_key])
                if chunk_key not in graph_record.mem_overhead_keys:
                    self._mem_quota_ref.release_quotas(
                        [build_quota_key(session_id, chunk_key, owner=graph_key)], _tell=True, _wait=False)
                if graph_record.preferred_data_device not in locations:
                    return storage_client.copy_to(session_id, [chunk_key], [graph_record.preferred_data_device])

        @log_unhandled
        def _handle_network_error(*exc):
            try:
                logger.warning('Communicating to %s encountered %s when fetching %s',
                               remote_addr, exc[0].__name__, chunk_key)

                if not storage_client.get_data_locations(session_id, [chunk_key])[0]:
                    logger.debug('Deleting chunk %s from worker because of failure of transfer', chunk_key)
                    storage_client.delete(session_id, [chunk_key])
                else:
                    # as data already transferred, we can skip the error
                    logger.debug('Chunk %s already transferred', chunk_key)
                    _finish_fetch()
                    return

                raise exc[1].with_traceback(exc[2])
            except (BrokenPipeError, ConnectionRefusedError, WorkerDead):
                self._resource_ref.detach_dead_workers([remote_addr], _tell=True, _wait=False)
                raise DependencyMissing((session_id, chunk_key))

        @log_unhandled
        def _fetch_step(sender_uid):
            if remote_addr in self._peer_blacklist:
                raise DependencyMissing((session_id, chunk_key))

            if graph_record.stop_requested:
                remote_disp_ref.register_free_slot(sender_uid, 'sender', _tell=True)
                raise ExecutionInterrupted

            sender_ref = self.promise_ref(sender_uid, address=remote_addr)
            logger.debug('Fetching chunk %s in %s to %s with slot %s',
                         chunk_key, remote_addr, self.address, sender_uid)

            return sender_ref.send_data(
                session_id, [chunk_key], [self.address], ensure_cached=ensure_cached,
                pin_token=graph_key, timeout=timeout, _timeout=timeout, _promise=True
            ).then(_finish_fetch)

        return promise.finished() \
            .then(lambda *_: remote_disp_ref.acquire_free_slot('sender', _promise=True)) \
            .then(_fetch_step) \
            .catch(_handle_network_error)

    def estimate_graph_finish_time(self, session_id, graph_key, calc_fetch=True, base_time=None):
        """
        Calc predictions for given chunk graph
        """
        session_graph_key = (session_id, graph_key)
        if session_graph_key not in self._graph_records:
            return
        graph_record = self._graph_records[session_graph_key]
        graph = graph_record.graph

        ops = set(type(c.op).__name__ for c in graph if not isinstance(c.op, Fetch))
        op_calc_key = ('calc_speed.' + list(ops)[0]) if len(ops) == 1 else None

        stats = defaultdict(lambda: dict(count=0))
        if self._status_ref:
            stats.update(self._status_ref.get_stats(['disk_read_speed', 'disk_write_speed',
                                                     'net_transfer_speed', op_calc_key]))

        if op_calc_key not in stats:
            return None
        if stats[op_calc_key]['count'] < options.optimize.min_stats_count:
            return None
        if abs(stats[op_calc_key]['count']) < 1e-6:
            return None

        input_size = 0
        net_size = 0
        disk_size = 0
        base_time = base_time or time.time()

        if calc_fetch:
            for c in graph:
                if not isinstance(c.op, Fetch):
                    break
                try:
                    data_size = calc_data_size(c)
                except (AttributeError, TypeError, ValueError):
                    data_size = 0
                input_size += data_size
                data_locations = self.storage_client.get_data_locations(session_id, [c.key])[0]
                if (0, DataStorageDevice.VINEYARD) in data_locations or \
                        (0, DataStorageDevice.SHARED_MEMORY) in data_locations:  # pragma: no cover
                    continue
                elif (0, DataStorageDevice.DISK) in data_locations:
                    disk_size += data_size
                else:
                    net_size += data_size

            if stats['net_transfer_speed']['count'] >= options.optimize.min_stats_count:
                base_time += net_size * 1.0 / stats['net_transfer_speed']['mean']
            if stats['disk_read_speed']['count'] >= options.optimize.min_stats_count:
                base_time += disk_size * 1.0 / stats['disk_read_speed']['mean']
            else:
                base_time += disk_size * 1.0 / options.optimize.default_disk_io_speed

        est_finish_time = base_time + input_size * 1.0 / stats[op_calc_key]['mean']

        graph_record.est_finish_time = est_finish_time
        self._status_ref.update_stats(dict(
            min_est_finish_time=min(rec.est_finish_time for rec in self._graph_records.values()),
            max_est_finish_time=max(rec.est_finish_time for rec in self._graph_records.values()),
        ), _tell=True, _wait=False)

        self.ref().estimate_graph_finish_time(session_id, graph_key, _tell=True, _delay=1)

    def _update_state(self, session_id, key, state):
        logger.debug('Operand %s switched to %s', key, getattr(state, 'name'))
        record = self._graph_records[(session_id, key)]
        record.state = state
        if self._status_ref:
            self._status_ref.update_progress(
                session_id, key, record.op_string, state, _tell=True, _wait=False)

    @promise.reject_on_exception
    @log_unhandled
    def execute_graph(self, session_id, graph_key, graph_ser, io_meta, data_metas,
                      calc_device=None, send_addresses=None, callback=None):
        """
        Submit graph to the worker and control the execution
        :param session_id: session id
        :param graph_key: graph key
        :param graph_ser: serialized executable graph
        :param io_meta: io meta of the chunk
        :param data_metas: data meta of each input chunk, as a dict
        :param calc_device: device for calculation, can be 'gpu' or 'cpu'
        :param send_addresses: targets to send results after execution
        :param callback: promise callback
        """
        session_graph_key = (session_id, graph_key)
        callback = callback or []
        if not isinstance(callback, list):
            callback = [callback]
        try:
            all_callbacks = self._graph_records[session_graph_key].finish_callbacks or []
            self._graph_records[session_graph_key].finish_callbacks.extend(callback)
            if not self._graph_records[session_graph_key].retry_pending:
                self._graph_records[session_graph_key].finish_callbacks = all_callbacks + callback
                return
        except KeyError:
            all_callbacks = []
        all_callbacks.extend(callback)

        calc_device = calc_device or 'cpu'
        if calc_device == 'cpu':  # pragma: no cover
            if options.vineyard.enabled:
                preferred_data_device = DataStorageDevice.VINEYARD
            else:
                preferred_data_device = DataStorageDevice.SHARED_MEMORY
        else:
            preferred_data_device = DataStorageDevice.CUDA

        # todo change this when handling multiple devices
        if preferred_data_device == DataStorageDevice.CUDA:
            slot = self._dispatch_ref.get_slots(calc_device)[0]
            proc_id = self.ctx.distributor.distribute(slot)
            preferred_data_device = (proc_id, preferred_data_device)

        graph_record = self._graph_records[(session_id, graph_key)] = GraphExecutionRecord(
            graph_ser, ExecutionState.ALLOCATING,
            io_meta=io_meta,
            data_metas=data_metas,
            chunk_targets=io_meta['chunks'],
            data_targets=list(io_meta.get('data_targets') or io_meta['chunks']),
            shared_input_chunks=set(io_meta.get('shared_input_chunks', [])),
            send_addresses=send_addresses,
            finish_callbacks=all_callbacks,
            calc_device=calc_device,
            preferred_data_device=preferred_data_device,
            pure_dep_chunk_keys=io_meta.get('pure_dep_chunk_keys') or set(),
        )

        _, long_op_string = concat_operand_keys(graph_record.graph, decompose=True)
        if long_op_string != graph_record.op_string:
            long_op_string = graph_record.op_string + ':' + long_op_string
        logger.debug('Worker graph %s(%s) targeting at %r accepted.', graph_key,
                     long_op_string, graph_record.chunk_targets)
        self._update_state(session_id, graph_key, ExecutionState.ALLOCATING)

        try:
            del self._result_cache[session_graph_key]
        except KeyError:
            pass

        @log_unhandled
        def _handle_success(*_):
            self._invoke_finish_callbacks(session_id, graph_key)

        @log_unhandled
        def _handle_rejection(*exc):
            # some error occurred...
            logger.debug('Entering _handle_rejection() for graph %s', graph_key)
            self._dump_execution_states()

            if graph_record.stop_requested:
                graph_record.stop_requested = False
                if not isinstance(exc[1], ExecutionInterrupted):
                    exc = build_exc_info(ExecutionInterrupted)

            if isinstance(exc[1], ExecutionInterrupted):
                logger.warning('Execution of graph %s interrupted.', graph_key)
            else:
                logger.exception('Unexpected error occurred in executing graph %s', graph_key, exc_info=exc)

            self._result_cache[(session_id, graph_key)] = GraphResultRecord(*exc, succeeded=False)
            self._invoke_finish_callbacks(session_id, graph_key)

        # collect target data already computed
        attrs = self.storage_client.get_data_attrs(session_id, graph_record.data_targets)
        save_attrs = dict((k, v) for k, v in zip(graph_record.data_targets, attrs) if v)

        # when all target data are computed, report success directly
        if all(k in save_attrs for k in graph_record.data_targets):
            logger.debug('All predecessors of graph %s already computed, call finish directly.', graph_key)
            sizes = dict((k, v.size) for k, v in save_attrs.items())
            shapes = dict((k, v.shape) for k, v in save_attrs.items())
            self._result_cache[(session_id, graph_key)] = GraphResultRecord(sizes, shapes)
            _handle_success()
        else:
            try:
                quota_request = self._prepare_quota_request(session_id, graph_key)
            except PinDataKeyFailed:
                logger.debug('Failed to pin chunk for graph %s', graph_key)

                # cannot pin input chunks: retry later
                retry_delay = graph_record.retry_delay + 0.5 + random.random()
                graph_record.retry_delay = min(1 + graph_record.retry_delay, 30)
                graph_record.retry_pending = True

                self.ref().execute_graph(
                    session_id, graph_key, graph_record.graph_serialized, graph_record.io_meta,
                    graph_record.data_metas, calc_device=calc_device, send_addresses=send_addresses,
                    _tell=True, _delay=retry_delay)
                return

            promise.finished() \
                .then(lambda *_: self._mem_quota_ref.request_batch_quota(
                    quota_request, _promise=True) if quota_request else None) \
                .then(lambda *_: self._prepare_graph_inputs(session_id, graph_key)) \
                .then(lambda *_: self._dispatch_ref.acquire_free_slot(calc_device, _promise=True)) \
                .then(lambda uid: self._send_calc_request(session_id, graph_key, uid)) \
                .then(lambda saved_keys: self._store_results(session_id, graph_key, saved_keys)) \
                .then(_handle_success, _handle_rejection)

    def _release_shared_store_quotas(self, session_id, graph_key, keys):
        graph_record = self._graph_records[(session_id, graph_key)]
        quota_keys_to_release = [
            build_quota_key(session_id, k, owner=graph_key) for k in keys
            if k not in graph_record.mem_overhead_keys
        ]
        for quota_key in quota_keys_to_release:
            graph_record.mem_request.pop(quota_key, None)
        self._mem_quota_ref.release_quotas(quota_keys_to_release, _tell=True)

    @log_unhandled
    def _prepare_graph_inputs(self, session_id, graph_key):
        """
        Load input data from spilled storage and other workers
        :param session_id: session id
        :param graph_key: key of the execution graph
        """
        storage_client = self.storage_client
        graph_record = self._graph_records[(session_id, graph_key)]
        graph = graph_record.graph
        input_metas = graph_record.io_meta.get('input_data_metas', dict())

        if graph_record.stop_requested:
            raise ExecutionInterrupted

        logger.debug('Start preparing input data for graph %s', graph_key)
        self._update_state(session_id, graph_key, ExecutionState.PREPARING_INPUTS)
        prepare_promises = []

        input_keys = set()
        reducer_indexes = set()
        for chunk in graph:
            op = chunk.op
            if isinstance(op, Fetch):
                if chunk.key in graph_record.pure_dep_chunk_keys:
                    continue
                input_keys.add(op.source_key or chunk.key)
            elif isinstance(op, FetchShuffle):
                reducer_index = get_chunk_reducer_index(graph.successors(chunk)[0])
                for input_key in op.source_keys:
                    part_key = (input_key, reducer_index)
                    input_keys.add(part_key)
                    reducer_indexes.add(part_key)

        local_keys = graph_record.pinned_keys & set(input_keys)
        non_local_keys = [k for k in input_keys if k not in local_keys]
        non_local_locations = storage_client.get_data_locations(session_id, non_local_keys)
        copy_keys = set(k for k, loc in zip(non_local_keys, non_local_locations) if loc)
        remote_keys = [k for k in non_local_keys if k not in copy_keys]

        # handle local keys
        self._release_shared_store_quotas(session_id, graph_key, local_keys)
        # handle move keys
        prepare_promises.extend(
            self._prepare_copy_keys(session_id, graph_key, copy_keys))
        # handle remote keys
        prepare_promises.extend(
            self._prepare_remote_keys(session_id, graph_key, remote_keys, input_metas))

        logger.debug('Graph key %s: Targets %r, loaded keys %r, copy keys %s, remote keys %r',
                     graph_key, graph_record.chunk_targets, local_keys, copy_keys, remote_keys)
        p = promise.all_(prepare_promises) \
            .then(lambda *_: logger.debug('Data preparation for graph %s finished', graph_key))
        return p

    @log_unhandled
    def _prepare_copy_keys(self, session_id, graph_key, copy_keys):
        promises = []
        graph_record = self._graph_records[(session_id, graph_key)]
        ensure_shared_keys = [k for k in copy_keys if k in graph_record.shared_input_chunks]
        better_shared_keys = [k for k in copy_keys if k not in graph_record.shared_input_chunks
                              and not isinstance(k, tuple)]

        def _release_copied_keys(keys):
            actual_moved_keys = self._pin_shared_data_keys(session_id, graph_key, keys)
            self._release_shared_store_quotas(session_id, graph_key, actual_moved_keys)

        if ensure_shared_keys:
            promises.append(
                self.storage_client.copy_to(
                    session_id, ensure_shared_keys, [graph_record.preferred_data_device],
                    ensure=True, pin_token=graph_key)
                .then(lambda *_: _release_copied_keys(ensure_shared_keys),
                      lambda *_: _release_copied_keys(ensure_shared_keys))
            )
        if better_shared_keys:
            promises.append(
                self.storage_client.copy_to(
                    session_id, better_shared_keys, [graph_record.preferred_data_device],
                    ensure=False, pin_token=graph_key)
                .then(lambda *_: _release_copied_keys(better_shared_keys),
                      lambda *_: _release_copied_keys(better_shared_keys))
            )
        return promises

    @log_unhandled
    def _prepare_remote_keys(self, session_id, graph_key, remote_keys, input_metas):
        promises = []
        graph_record = self._graph_records[(session_id, graph_key)]

        filtered_remote_keys = remote_keys
        if self._receiver_manager_ref:
            transferring_keys = set(self._receiver_manager_ref.filter_receiving_keys(
                session_id, remote_keys))
            if transferring_keys:
                logger.debug('Data %s already scheduled for fetch, just wait.', transferring_keys)
                filtered_remote_keys = [k for k in remote_keys if k not in transferring_keys]
                promises.append(self._receiver_manager_ref.add_keys_callback(
                    session_id, transferring_keys, _promise=True,
                    _timeout=options.worker.prepare_data_timeout))

        for input_key in filtered_remote_keys:
            # load data from another worker
            chunk_meta = input_metas.get(input_key) or graph_record.data_metas.get(input_key)
            if chunk_meta:
                chunk_meta.workers = tuple(w for w in chunk_meta.workers if w != self.address)

            if chunk_meta is None or not chunk_meta.workers:
                raise DependencyMissing(f'Dependency {input_key} not met on sending')
            worker_results = list(chunk_meta.workers)
            random.shuffle(worker_results)

            # fetch data from other workers, if one fails, try another
            p = promise.finished(_accept=False)
            for worker in worker_results:
                p = p.catch(functools.partial(
                    self._fetch_remote_data, session_id, graph_key, input_key, worker,
                    ensure_cached=input_key in graph_record.shared_input_chunks))
            promises.append(p)
        return promises

    @log_unhandled
    def _send_calc_request(self, session_id, graph_key, calc_uid):
        """
        Start actual calculation in CpuCalcActor
        :param session_id: session id
        :param graph_key: key of the execution graph
        :param calc_uid: uid of the allocated CpuCalcActor
        """
        try:
            graph_record = self._graph_records[(session_id, graph_key)]

            if graph_record.stop_requested:
                raise ExecutionInterrupted

            graph_record.calc_actor_uid = calc_uid

            self._update_state(session_id, graph_key, ExecutionState.CALCULATING)
            calc_ref = self.promise_ref(calc_uid)

            def _start_calc(*_):
                logger.debug('Submit calculation for graph %s in actor %s', graph_key, calc_uid)
                if self._daemon_ref is None or self._daemon_ref.is_actor_process_alive(calc_ref):
                    return calc_ref.calc(
                        session_id, graph_key, graph_record.graph_serialized, graph_record.chunk_targets,
                        _promise=True
                    )
                else:
                    raise WorkerProcessStopped

            self.estimate_graph_finish_time(session_id, graph_key, calc_fetch=False)
        except:  # noqa: E722
            self._dispatch_ref.register_free_slot(calc_uid, 'cpu', _tell=True)
            raise

        # make sure that memory suffices before actually run execution
        mem_request = graph_record.mem_request
        if mem_request:
            logger.debug('Ensuring resource %r for graph %s', mem_request, graph_key)
            return self._mem_quota_ref.request_batch_quota(mem_request, process_quota=True, _promise=True) \
                .then(_start_calc)
        else:
            return _start_calc()

    @log_unhandled
    def _do_active_transfer(self, session_id, graph_key, data_to_addresses):
        if graph_key:
            logger.debug('Start active transfer for graph %s', graph_key)
            graph_record = self._graph_records[session_id, graph_key]
        else:
            graph_record = None

        # transfer the result chunk to expected endpoints
        @log_unhandled
        def _send_chunk(sender_uid, data_keys, target_addr):
            if graph_record and graph_record.stop_requested:
                self._dispatch_ref.register_free_slot(sender_uid, 'sender', _tell=True)
                raise ExecutionInterrupted

            sender_ref = self.promise_ref(sender_uid)
            timeout = options.worker.prepare_data_timeout
            logger.debug('Actively sending chunks %r to %s with slot %s',
                         data_keys, target_addr, sender_uid)
            return sender_ref.send_data(session_id, data_keys, [target_addr], ensure_cached=False,
                                        timeout=timeout, _timeout=timeout, _promise=True)

        if graph_record:
            if graph_record.mem_request:
                self._mem_quota_ref.release_quotas(tuple(graph_record.mem_request.keys()), _tell=True)
                graph_record.mem_request = dict()

        promises = []
        address_to_data_keys = defaultdict(set)
        for key, targets in data_to_addresses.items():
            for target in targets:
                address_to_data_keys[target].add(key)
        for target, keys in address_to_data_keys.items():
            if target == self.address:
                continue
            promises.append(promise.finished()
                            .then(lambda: self._dispatch_ref.acquire_free_slot('sender', _promise=True))
                            .then(functools.partial(_send_chunk, data_keys=keys, target_addr=target))
                            .catch(lambda *_: None))

        for target, keys in address_to_data_keys.items():
            self.ctx.actor_ref(ReceiverManagerActor.default_uid(), address=target) \
                .register_pending_keys(session_id, keys, _tell=True, _wait=False)

        return promise.all_(promises)

    @log_unhandled
    def _store_results(self, session_id, graph_key, saved_keys):
        """
        Store calc results into shared cache or spill
        :param session_id: session id
        :param graph_key: key of the execution graph
        """
        storage_client = self.storage_client

        self.deallocate_scheduler_resource(session_id, graph_key)

        graph_record = self._graph_records[session_id, graph_key]
        chunk_keys = graph_record.chunk_targets
        send_addresses = graph_record.send_addresses

        logger.debug('Graph %s: Start putting %r into shared cache using actor uid %s.',
                     graph_key, chunk_keys, graph_record.calc_actor_uid)
        self._update_state(session_id, graph_key, ExecutionState.STORING)

        raw_calc_ref = self.ctx.actor_ref(graph_record.calc_actor_uid)
        calc_ref = self.promise_ref(raw_calc_ref)

        if graph_record.stop_requested:
            logger.debug('Graph %s already marked for stop, quit.', graph_key)
            if (self._daemon_ref is None or self._daemon_ref.is_actor_process_alive(raw_calc_ref)) \
                    and self.ctx.has_actor(raw_calc_ref):
                logger.debug('Try remove keys for graph %s.', graph_key)
                raw_calc_ref.remove_cache(session_id, list(saved_keys), _tell=True)
            logger.debug('Graph %s already marked for stop, quit.', graph_key)
            raise ExecutionInterrupted

        storage_client.unpin_data_keys(session_id, graph_record.pinned_keys, graph_key)
        self._dump_execution_states()

        data_attrs = storage_client.get_data_attrs(session_id, saved_keys)

        if self._daemon_ref is not None and not self._daemon_ref.is_actor_process_alive(raw_calc_ref):
            raise WorkerProcessStopped

        def _cache_result(*_):
            save_sizes = dict((k, v.size) for k, v in zip(saved_keys, data_attrs) if v)
            save_shapes = dict((k, v.shape) for k, v in zip(saved_keys, data_attrs) if v)
            self._result_cache[(session_id, graph_key)] = GraphResultRecord(save_sizes, save_shapes)

        if not send_addresses:
            # no endpoints to send, dump keys into shared memory and return
            logger.debug('Worker graph %s(%s) finished execution. Dumping results...',
                         graph_key, graph_record.op_string)
            return calc_ref.store_results(session_id, graph_key, saved_keys, data_attrs, _promise=True) \
                .then(_cache_result)
        else:
            # dump keys into shared memory and send
            all_addresses = [{v} if isinstance(v, str) else set(v)
                             for v in send_addresses.values()]
            all_addresses = list(functools.reduce(lambda a, b: a | b, all_addresses, set()))
            logger.debug('Worker graph %s(%s) finished execution. Dumping results '
                         'while actively transferring into %r...',
                         graph_key, graph_record.op_string, all_addresses)

            data_to_addresses = dict((k, v) for k, v in send_addresses.items()
                                     if k in saved_keys)

            return calc_ref.store_results(session_id, graph_key, saved_keys, data_attrs, _promise=True) \
                .then(_cache_result) \
                .then(lambda *_: functools.partial(self._do_active_transfer,
                                                   session_id, graph_key, data_to_addresses))

    def deallocate_scheduler_resource(self, session_id, graph_key, delay=0):
        try:
            graph_record = self._graph_records[(session_id, graph_key)]
            if not graph_record.resource_released:
                graph_record.resource_released = True
                self._resource_ref.deallocate_resource(
                    session_id, graph_key, self.address, _delay=delay, _tell=True, _wait=False)
        except:  # noqa: E722
            pass

    def _cleanup_graph(self, session_id, graph_key):
        """
        Do clean up after graph is executed
        :param session_id: session id
        :param graph_key: graph key
        """
        logger.debug('Cleaning callbacks for graph %s', graph_key)
        try:
            graph_record = self._graph_records[(session_id, graph_key)]
        except KeyError:
            return

        self.deallocate_scheduler_resource(session_id, graph_key)

        mem_quota_keys = tuple(graph_record.mem_request.keys())
        self._mem_quota_ref.cancel_requests(mem_quota_keys, _tell=True)
        if graph_record.mem_request:
            self._mem_quota_ref.release_quotas(mem_quota_keys, _tell=True)
            if graph_record.calc_actor_uid:
                target_proc_id = self.ctx.distributor.distribute(graph_record.calc_actor_uid)
                owned_quota_keys = tuple(change_quota_key_owner(k, target_proc_id)
                                         for k in mem_quota_keys)
                self._mem_quota_ref.release_quotas(owned_quota_keys, _tell=True)

        if graph_record.pinned_keys:
            self.storage_client.unpin_data_keys(session_id, graph_record.pinned_keys, graph_key)

        if self._status_ref:
            self._status_ref.remove_progress(session_id, graph_key, _tell=True, _wait=False)
        del self._graph_records[(session_id, graph_key)]
        self.ref().delete_result_cache(session_id, graph_key, _tell=True, _delay=20)

    @promise.reject_on_exception
    @log_unhandled
    def add_finish_callback(self, session_id, graph_key, callback):
        """
        Register a callback to callback store
        :param session_id: session id
        :param graph_key: graph key
        :param callback: promise call
        """
        logger.debug('Adding callback %r for graph %s', callback, graph_key)
        try:
            args, kwargs = self._result_cache[(session_id, graph_key)].build_args()
            self.tell_promise(callback, *args, **kwargs)
        except KeyError:
            self._graph_records[(session_id, graph_key)].finish_callbacks.append(callback)

    @promise.reject_on_exception
    @log_unhandled
    def send_data_to_workers(self, session_id, data_to_addresses, callback=None):
        """
        Send data stored in worker to other addresses
        :param session_id: session id
        :param data_to_addresses: dict mapping data to a list of addresses
        :param callback: promise call
        :return:
        """
        self._do_active_transfer(session_id, None, data_to_addresses) \
            .then(lambda *_: self.tell_promise(callback) if callback else None) \
            .catch(lambda *exc: self.tell_promise(callback, *exc, _accept=False) if callback else None)

    @log_unhandled
    def stop_execution(self, session_id, graph_key):
        """
        Mark graph for stopping
        :param session_id: session id
        :param graph_key: graph key
        """
        logger.debug('Receive stop for graph %s', graph_key)
        try:
            graph_record = self._graph_records[(session_id, graph_key)]
        except KeyError:
            return

        graph_record.stop_requested = True
        if graph_record.state == ExecutionState.ALLOCATING:
            if graph_record.mem_request:
                self._mem_quota_ref.cancel_requests(
                    tuple(graph_record.mem_request.keys()), build_exc_info(ExecutionInterrupted), _tell=True)
        elif graph_record.state == ExecutionState.CALCULATING:
            if options.worker.recover_dead_process \
                    and self._daemon_ref is not None and graph_record.calc_actor_uid is not None:
                self._daemon_ref.kill_actor_process(self.ctx.actor_ref(graph_record.calc_actor_uid), _tell=True)

    @log_unhandled
    def delete_data_by_keys(self, session_id, keys):
        self.storage_client.delete(session_id, keys, _tell=True)

    @log_unhandled
    def handle_worker_change(self, _adds, removes):
        """
        Handle worker dead event
        :param removes: list of dead workers
        """
        if not removes:
            return
        self._peer_blacklist.update(removes)

        handled_refs = self.reject_dead_endpoints(removes, *build_exc_info(DependencyMissing))
        logger.debug('Peer worker halt received. Affected promises %r rejected.',
                     [ref.uid for ref in handled_refs])

        for sender_slot in self._dispatch_ref.get_slots('sender'):
            self.ctx.actor_ref(sender_slot).reject_dead_endpoints(removes, _tell=True)
        for receiver_slot in self._dispatch_ref.get_slots('receiver'):
            self.ctx.actor_ref(receiver_slot).notify_dead_senders(removes, _tell=True)

    @log_unhandled
    def _invoke_finish_callbacks(self, session_id, graph_key):
        """
        Call finish callback when execution is done
        :param session_id: session id
        :param graph_key: graph key
        """
        query_key = (session_id, graph_key)
        callbacks = self._graph_records[query_key].finish_callbacks
        args, kwargs = self._result_cache[query_key].build_args()

        logger.debug('Send finish callback for graph %s into %d targets', graph_key, len(callbacks))
        kwargs['_wait'] = False
        for cb in callbacks:
            self.tell_promise(cb, *args, **kwargs)
        self._cleanup_graph(session_id, graph_key)

    def delete_result_cache(self, session_id, graph_key):
        del self._result_cache[(session_id, graph_key)]

    def _dump_execution_states(self):
        if logger.getEffectiveLevel() <= logging.DEBUG:
            cur_time = time.time()
            states = dict((k[1], (cur_time - v.state_time, v.state.name))
                          for k, v in self._graph_records.items()
                          if v.state != ExecutionState.ALLOCATING)
            logger.debug('Executing states: %r', states)
