# Buggy code
def _replace_copied_tilebale(self, graph):
    if len(self._optimizer_context) == 0:
        return graph

    new_graph = DAG()
    replaced_tileables = weakref.WeakKeyDictionary()
    for n in graph.topological_iter():
        if graph.count_predecessors(n) == 0:
            if n in self._optimizer_context and all(
                suc in self._optimizer_context for suc in graph.successors(n)
            ):
                replaced_tileables[n] = new_node = self._optimizer_context[n]
            else:
                new_node = n
        elif any(inp in replaced_tileables for inp in n.inputs):
            new_inputs = [replaced_tileables.get(i, i) for i in n.inputs]
            new_tileables = copy_tileables(n.op.outputs, inputs=new_inputs)
            for t, new_t in zip(n.op.outputs, new_tileables):
                replaced_tileables[t] = new_t.data
                if t is n:
                    new_node = new_t.data
        else:
            new_node = n
        new_graph.add_node(new_node)
        for inp in new_node.inputs:
            new_graph.add_node(inp)
            new_graph.add_edge(inp, new_node)
    self._optimizer_context.update(replaced_tileables)
    return new_graph


# Traceback (KeyError)
In [9]: d = md.read_csv('Downloads/test.csv')

In [10]: (d.head() + 1).execute()
Unexpected exception occurred in GraphActor.prepare_graph.
Traceback (most recent call last):
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/utils.py", line 383, in _wrapped
return func(*args, **kwargs)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/utils.py", line 457, in _inner
return func(*args, **kwargs)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/scheduler/graph.py", line 666, in prepare_graph
self._target_tileable_datas + fetch_tileables, tileable_graph)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/utils.py", line 457, in _inner
return func(*args, **kwargs)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/tiles.py", line 348, in build
tileables, tileable_graph=tileable_graph)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/utils.py", line 457, in _inner
return func(*args, **kwargs)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/tiles.py", line 262, in build
self._on_tile_failure(tileable_data.op, exc_info)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/tiles.py", line 301, in inner
raise exc_info[1].with_traceback(exc_info[2]) from None
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/tiles.py", line 242, in build
tiled = self._tile(tileable_data, tileable_graph)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/tiles.py", line 337, in _tile
return super()._tile(tileable_data, tileable_graph)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/tiles.py", line 191, in _tile
inputs=[cache[inp] for inp in tileable_data.inputs],
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/tiles.py", line 191, in <listcomp>
inputs=[cache[inp] for inp in tileable_data.inputs],
File "/Users/hekaisheng/miniconda3/envs/py3.7.2/lib/python3.7/weakref.py", line 394, in __getitem__
return self.data[ref(key)]
KeyError: <weakref at 0x192f4965e8; to 'DataFrameData' at 0x192f4a2048>
Failed to start graph execution.
Traceback (most recent call last):
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/scheduler/graph.py", line 411, in _execute_graph
self.prepare_graph(compose=compose)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/utils.py", line 383, in _wrapped
return func(*args, **kwargs)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/utils.py", line 457, in _inner
return func(*args, **kwargs)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/scheduler/graph.py", line 666, in prepare_graph
self._target_tileable_datas + fetch_tileables, tileable_graph)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/utils.py", line 457, in _inner
return func(*args, **kwargs)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/tiles.py", line 348, in build
tileables, tileable_graph=tileable_graph)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/utils.py", line 457, in _inner
return func(*args, **kwargs)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/tiles.py", line 262, in build
self._on_tile_failure(tileable_data.op, exc_info)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/tiles.py", line 301, in inner
raise exc_info[1].with_traceback(exc_info[2]) from None
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/tiles.py", line 242, in build
tiled = self._tile(tileable_data, tileable_graph)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/tiles.py", line 337, in _tile
return super()._tile(tileable_data, tileable_graph)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/tiles.py", line 191, in _tile
inputs=[cache[inp] for inp in tileable_data.inputs],
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/tiles.py", line 191, in <listcomp>
inputs=[cache[inp] for inp in tileable_data.inputs],
File "/Users/hekaisheng/miniconda3/envs/py3.7.2/lib/python3.7/weakref.py", line 394, in __getitem__
return self.data[ref(key)]
KeyError: <weakref at 0x192f4965e8; to 'DataFrameData' at 0x192f4a2048>
Unexpected exception occurred in GraphActor.get_chunk_graph.
Traceback (most recent call last):
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/utils.py", line 383, in _wrapped
return func(*args, **kwargs)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/scheduler/graph.py", line 471, in get_chunk_graph
raise GraphNotExists from None
mars.errors.GraphNotExists
Unexpected exception occurred in GraphActor.execute_graph.
Traceback (most recent call last):
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/utils.py", line 383, in _wrapped
return func(*args, **kwargs)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/scheduler/graph.py", line 383, in execute_graph
self._execute_graph(compose=compose)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/scheduler/graph.py", line 411, in _execute_graph
self.prepare_graph(compose=compose)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/utils.py", line 383, in _wrapped
return func(*args, **kwargs)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/utils.py", line 457, in _inner
return func(*args, **kwargs)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/scheduler/graph.py", line 666, in prepare_graph
self._target_tileable_datas + fetch_tileables, tileable_graph)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/utils.py", line 457, in _inner
return func(*args, **kwargs)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/tiles.py", line 348, in build
tileables, tileable_graph=tileable_graph)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/utils.py", line 457, in _inner
return func(*args, **kwargs)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/tiles.py", line 262, in build
self._on_tile_failure(tileable_data.op, exc_info)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/tiles.py", line 301, in inner
raise exc_info[1].with_traceback(exc_info[2]) from None
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/tiles.py", line 242, in build
tiled = self._tile(tileable_data, tileable_graph)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/tiles.py", line 337, in _tile
return super()._tile(tileable_data, tileable_graph)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/tiles.py", line 191, in _tile
inputs=[cache[inp] for inp in tileable_data.inputs],
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/tiles.py", line 191, in <listcomp>
inputs=[cache[inp] for inp in tileable_data.inputs],
File "/Users/hekaisheng/miniconda3/envs/py3.7.2/lib/python3.7/weakref.py", line 394, in __getitem__
return self.data[ref(key)]
KeyError: <weakref at 0x192f4965e8; to 'DataFrameData' at 0x192f4a2048>
Traceback (most recent call last):
File "src/gevent/greenlet.py", line 854, in gevent._gevent_cgreenlet.Greenlet.run
File "mars/actors/pool/gevent_pool.pyx", line 90, in mars.actors.pool.gevent_pool.ActorExecutionContext.fire_run
File "mars/actors/pool/gevent_pool.pyx", line 93, in mars.actors.pool.gevent_pool.ActorExecutionContext.fire_run
File "mars/actors/pool/gevent_pool.pyx", line 104, in mars.actors.pool.gevent_pool.ActorExecutionContext.fire_run
File "mars/actors/pool/gevent_pool.pyx", line 98, in mars.actors.pool.gevent_pool.ActorExecutionContext.fire_run
File "mars/actors/core.pyx", line 112, in mars.actors.core._FunctionActor.on_receive
File "mars/actors/core.pyx", line 114, in mars.actors.core._FunctionActor.on_receive
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/utils.py", line 383, in _wrapped
return func(*args, **kwargs)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/scheduler/graph.py", line 383, in execute_graph
self._execute_graph(compose=compose)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/scheduler/graph.py", line 411, in _execute_graph
self.prepare_graph(compose=compose)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/utils.py", line 383, in _wrapped
return func(*args, **kwargs)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/utils.py", line 457, in _inner
return func(*args, **kwargs)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/scheduler/graph.py", line 666, in prepare_graph
self._target_tileable_datas + fetch_tileables, tileable_graph)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/utils.py", line 457, in _inner
return func(*args, **kwargs)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/tiles.py", line 348, in build
tileables, tileable_graph=tileable_graph)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/utils.py", line 457, in _inner
return func(*args, **kwargs)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/tiles.py", line 262, in build
self._on_tile_failure(tileable_data.op, exc_info)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/tiles.py", line 301, in inner
raise exc_info[1].with_traceback(exc_info[2]) from None
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/tiles.py", line 242, in build
tiled = self._tile(tileable_data, tileable_graph)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/tiles.py", line 337, in _tile
return super()._tile(tileable_data, tileable_graph)
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/tiles.py", line 191, in _tile
inputs=[cache[inp] for inp in tileable_data.inputs],
File "/Users/hekaisheng/Documents/mars_dev/mars/mars/tiles.py", line 191, in <listcomp>
inputs=[cache[inp] for inp in tileable_data.inputs],
File "/Users/hekaisheng/miniconda3/envs/py3.7.2/lib/python3.7/weakref.py", line 394, in __getitem__
return self.data[ref(key)]
KeyError: <weakref at 0x192f4965e8; to 'DataFrameData' at 0x192f4a2048>
2021-01-25T06:20:24Z <Greenlet at 0x192f402ae8: <built-in method fire_run of mars.actors.pool.gevent_pool.ActorExecutionContext object at 0x192f391598>> failed with KeyError

Traceback (most recent call last):
File "src/gevent/greenlet.py", line 854, in gevent._gevent_cgreenlet.Greenlet.run
File "mars/actors/pool/gevent_pool.pyx", line 548, in mars.actors.pool.gevent_pool.ActorRemoteHelper._send_remote
File "mars/actors/pool/gevent_pool.pyx", line 552, in mars.actors.pool.gevent_pool.ActorRemoteHelper._send_remote
File "mars/actors/pool/gevent_pool.pyx", line 559, in mars.actors.pool.gevent_pool.ActorRemoteHelper._send_remote
File "mars/actors/pool/gevent_pool.pyx", line 1036, in mars.actors.pool.gevent_pool.Communicator._on_receive_send
File "mars/actors/pool/messages.pyx", line 747, in mars.actors.pool.messages.pack_result_message
_pack_object(result, buf)
File "mars/actors/pool/messages.pyx", line 305, in mars.actors.pool.messages._pack_object
m = dumps(obj)
TypeError: can't pickle weakref objects
2021-01-25T06:20:24Z <Greenlet at 0x1932ad0e18: <built-in method _send_remote of mars.actors.pool.gevent_pool.ActorRemoteHelper object at 0x1932cf5ea8>('0.0.0.0:44189', [bytearray(b'\x05\x01 \x00\x00\x00\x00\x00\x00\x00)> failed with TypeError

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-10-9da2df47dde0> in <module>
----> 1 (d.head() + 1).execute()

~/Documents/mars_dev/mars/mars/core.py in execute(self, session, **kw)
644
645         if wait:
--> 646             return run()
647         else:
648             thread_executor = ThreadPoolExecutor(1)

~/Documents/mars_dev/mars/mars/core.py in run()
640
641         def run():
--> 642             self.data.execute(session, **kw)
643             return self
644

~/Documents/mars_dev/mars/mars/core.py in execute(self, session, **kw)
377
378         if wait:
--> 379             return run()
380         else:
381             # leverage ThreadPoolExecutor to submit task,

~/Documents/mars_dev/mars/mars/core.py in run()
372         def run():
373             # no more fetch, thus just fire run
--> 374             session.run(self, **kw)
375             # return Tileable or ExecutableTuple itself
376             return self

~/Documents/mars_dev/mars/mars/session.py in run(self, *tileables, **kw)
503         tileables = tuple(mt.tensor(t) if not isinstance(t, (Entity, Base)) else t
504                           for t in tileables)
--> 505         result = self._sess.run(*tileables, **kw)
506
507         for t in tileables:

~/Documents/mars_dev/mars/mars/session.py in run(self, *tileables, **kw)
316                     break
317                 if graph_state == GraphState.FAILED:
--> 318                     exc_info = self._api.get_graph_exc_info(self._session_id, graph_key)
319                     if exc_info is not None:
320                         exc = exc_info[1].with_traceback(exc_info[2])

~/Documents/mars_dev/mars/mars/api.py in get_graph_exc_info(self, session_id, graph_key)
193         graph_meta_ref = self.get_graph_meta_ref(session_id, graph_key)
194         try:
--> 195             return graph_meta_ref.get_exc_info()
196         except ActorNotExist:
197             raise GraphNotExists

~/Documents/mars_dev/mars/mars/actors/core.pyx in mars.actors.core.ActorRef.__getattr__._mt_call()

~/Documents/mars_dev/mars/mars/actors/core.pyx in mars.actors.core.ActorRef.send()

~/Documents/mars_dev/mars/mars/actors/pool/gevent_pool.pyx in mars.actors.pool.gevent_pool.ActorRemoteHelper.send()

~/Documents/mars_dev/mars/mars/actors/pool/gevent_pool.pyx in mars.actors.pool.gevent_pool.ActorRemoteHelper.send()

~/Documents/mars_dev/mars/mars/actors/pool/gevent_pool.pyx in mars.actors.pool.gevent_pool.ActorRemoteHelper._send()

~/miniconda3/envs/py3.7.2/lib/python3.7/site-packages/gevent/pool.py in apply(self, func, args, kwds)
159         if self._apply_immediately():
160             return func(*args, **kwds)
--> 161         return self.spawn(func, *args, **kwds).get()
162
163     def __map(self, func, iterable):

~/miniconda3/envs/py3.7.2/lib/python3.7/site-packages/gevent/_gevent_cgreenlet.cpython-37m-darwin.so in gevent._gevent_cgreenlet.Greenlet.get()

~/miniconda3/envs/py3.7.2/lib/python3.7/site-packages/gevent/_gevent_cgreenlet.cpython-37m-darwin.so in gevent._gevent_cgreenlet.Greenlet._raise_exception()

~/miniconda3/envs/py3.7.2/lib/python3.7/site-packages/gevent/_compat.py in reraise(t, value, tb)
63     def reraise(t, value, tb=None): # pylint:disable=unused-argument
64         if value.__traceback__ is not tb and tb is not None:
---> 65             raise value.with_traceback(tb)
66         raise value
67     def exc_clear():

~/miniconda3/envs/py3.7.2/lib/python3.7/site-packages/gevent/_gevent_cgreenlet.cpython-37m-darwin.so in gevent._gevent_cgreenlet.Greenlet.run()

~/Documents/mars_dev/mars/mars/actors/pool/gevent_pool.pyx in mars.actors.pool.gevent_pool.ActorRemoteHelper._send_remote()

~/Documents/mars_dev/mars/mars/actors/pool/gevent_pool.pyx in mars.actors.pool.gevent_pool.ActorRemoteHelper._send_remote()

~/Documents/mars_dev/mars/mars/actors/pool/gevent_pool.pyx in mars.actors.pool.gevent_pool.ActorRemoteHelper._send_remote()

~/Documents/mars_dev/mars/mars/actors/pool/gevent_pool.pyx in mars.actors.pool.gevent_pool.Communicator._on_receive_send()

~/Documents/mars_dev/mars/mars/actors/pool/messages.pyx in mars.actors.pool.messages.pack_result_message()
745     _pack_index(from_index, buf)
746     _pack_index(to_index, buf)
--> 747     _pack_object(result, buf)
748
749     if write is not None:

~/Documents/mars_dev/mars/mars/actors/pool/messages.pyx in mars.actors.pool.messages._pack_object()
303     else:
304         st = PICKLE
--> 305         m = dumps(obj)
306
307     _pack_byte(st, buf)

TypeError: can't pickle weakref objects

# Fixed code
def _replace_copied_tilebale(self, graph):
    if len(self._optimizer_context) == 0:
        return graph

    new_graph = DAG()
    replaced_tileables = weakref.WeakKeyDictionary()
    for n in graph.topological_iter():
        if graph.count_predecessors(n) == 0:
            if n in self._optimizer_context and all(
                suc in self._optimizer_context for suc in graph.successors(n)
            ):
                replaced_tileables[n] = new_node = self._optimizer_context[n]
            else:
                new_node = n
        elif any(inp in replaced_tileables for inp in n.inputs) or any(
            inp not in new_graph for inp in n.inputs
        ):
            new_inputs = []
            for i in n.inputs:
                if i in replaced_tileables:
                    new_inputs.append(replaced_tileables[i])
                elif i not in graph:
                    new_inputs.append(self._optimizer_context[i])
                else:
                    new_inputs.append(i)
            new_tileables = copy_tileables(n.op.outputs, inputs=new_inputs)
            for t, new_t in zip(n.op.outputs, new_tileables):
                replaced_tileables[t] = new_t.data
                if t is n:
                    new_node = new_t.data
        else:
            new_node = n
        new_graph.add_node(new_node)
        for inp in new_node.inputs:
            new_graph.add_node(inp)
            new_graph.add_edge(inp, new_node)
    self._optimizer_context.update(replaced_tileables)
    return new_graph