Repository: mars-project/mars

Issue:
[BUG] Execute df.head() + 1 failed in distributed mode
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**
Reproduce:

``` Python
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
```



Discussion:

Resolving patch:
diff --git a/mars/optimizes/tileable_graph/core.py b/mars/optimizes/tileable_graph/core.py
--- a/mars/optimizes/tileable_graph/core.py
+++ b/mars/optimizes/tileable_graph/core.py
@@ -98,8 +98,16 @@ def _replace_copied_tilebale(self, graph):
                     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]
+            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


Test patch:
diff --git a/mars/optimizes/tileable_graph/tests/test_head.py b/mars/optimizes/tileable_graph/tests/test_head.py
--- a/mars/optimizes/tileable_graph/tests/test_head.py
+++ b/mars/optimizes/tileable_graph/tests/test_head.py
@@ -71,6 +71,11 @@ def testReadCSVHead(self):
 
                 self.assertIn('cannot run iloc', str(cm.exception))
 
+            with self._raise_iloc():
+                s = mdf.head(5).sum()
+                expected = df.head(5).sum()
+                pd.testing.assert_series_equal(s.execute().fetch(), expected)
+
             pd.testing.assert_frame_equal(
                 mdf.head(99).execute().fetch().reset_index(drop=True), df.head(99))