Repository: mars-project/mars

Issue:
[BUG] Tensor does not return data when eager mode is on
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**

First, I execute a tensor with eager mode off, it succeeded, then I turn on eager mode, execute another tensor with some error happened, finally, repr the identical tensor to the first tensor does not trigger execution automatically.

**To Reproduce**

```
In [1]: import mars.tensor as mt                                                    

In [2]: from mars.config import options                                             

In [3]: a = mt.array( 
   ...:     [[0.1, 0.2, 0.3], 
   ...:      [0.3, 0.4, 0.2]] 
   ...:     )                                                                       

In [4]: (a[:, mt.newaxis, :] - a[mt.newaxis, ...]).execute()                        
Out[4]: 
array([[[ 0. ,  0. ,  0. ],
        [-0.2, -0.2,  0.1]],

       [[ 0.2,  0.2, -0.1],
        [ 0. ,  0. ,  0. ]]])

In [5]: options.eager_mode = True                                                   

In [6]:  mt.sqrt(mt.sum((a[:, mt.newaxis, :] - a[mt.newaxis, ...]) ** 2, axis=-1))  
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
~/Workspace/mars/mars/tiles.py in _dispatch(self, op)
    110         try:
--> 111             handler = self._handlers[op_cls]
    112             return handler(op)

KeyError: <class 'mars.tensor.expressions.arithmetic.sqrt.TensorSqrt'>

During handling of the above exception, another exception occurred:

ValueError                                Traceback (most recent call last)
<ipython-input-6-de975d9f2cf8> in <module>
----> 1 mt.sqrt(mt.sum((a[:, mt.newaxis, :] - a[mt.newaxis, ...]) ** 2, axis=-1))

~/Workspace/mars/mars/tensor/expressions/utils.py in h(*tensors, **kw)
    157                 kw['dtype'] = dtype
    158 
--> 159             ret = func(*tensors, **kw)
    160             if ret is NotImplemented:
    161                 reverse_func = getattr(inspect.getmodule(func), 'r{0}'.format(func.__name__), None) \

~/Workspace/mars/mars/tensor/expressions/arithmetic/sqrt.py in sqrt(x, out, where, **kwargs)
     78     """
     79     op = TensorSqrt(**kwargs)
---> 80     return op(x, out=out, where=where)
     81 

~/Workspace/mars/mars/tensor/expressions/arithmetic/core.py in __call__(self, x, out, where)
    348 
    349     def __call__(self, x, out=None, where=None):
--> 350         return self._call(x, out=out, where=where)
    351 
    352 

~/Workspace/mars/mars/tensor/expressions/arithmetic/core.py in _call(self, x, out, where)
    331         shape = x.shape
    332 
--> 333         t = self.new_tensor([x, out, where], shape)
    334 
    335         if out is None:

~/Workspace/mars/mars/tensor/expressions/core.py in new_tensor(self, inputs, shape, dtype, **kw)
     54             raise TypeError('cannot new tensor with more than 1 outputs')
     55 
---> 56         return self.new_tensors(inputs, shape, dtype=dtype, **kw)[0]
     57 
     58     def calc_shape(self, *inputs_shape):

~/Workspace/mars/mars/tensor/expressions/core.py in new_tensors(self, inputs, shape, dtype, chunks, nsplits, output_limit, kws, **kw)
     48                     output_limit=None, kws=None, **kw):
     49         return self.new_entities(inputs, shape, chunks=chunks, nsplits=nsplits,
---> 50                                  output_limit=output_limit, kws=kws, dtype=dtype, **kw)
     51 
     52     def new_tensor(self, inputs, shape, dtype=None, **kw):

~/Workspace/mars/mars/core.py in new_entities(self, inputs, shape, **kwargs)
    578         entities = self._new_entities(inputs, shape, **kwargs)
    579         if is_eager_mode():
--> 580             ExecutableTuple(entities).execute(fetch=False)
    581         return entities
    582 

~/Workspace/mars/mars/core.py in execute(self, session, **kw)
    594         if session is None:
    595             session = Session.default_or_local()
--> 596         return session.run(*self, **kw)

~/Workspace/mars/mars/session.py in run(self, *tensors, **kw)
    107 
    108         tensors = tuple(mt.tensor(t) for t in tensors)
--> 109         result = self._sess.run(*tensors, **kw)
    110 
    111         for t in tensors:

~/Workspace/mars/mars/session.py in run(self, *tensors, **kw)
     49         if 'n_parallel' not in kw:
     50             kw['n_parallel'] = cpu_count()
---> 51         res = self._executor.execute_tensors(tensors, **kw)
     52         return res
     53 

~/Workspace/mars/mars/utils.py in _wrapped(*args, **kwargs)
    352     def _wrapped(*args, **kwargs):
    353         _kernel_mode.eager = False
--> 354         return_value = func(*args, **kwargs)
    355         _kernel_mode.eager = None
    356         return return_value

~/Workspace/mars/mars/executor.py in execute_tensors(self, tensors, fetch, n_parallel, n_thread, print_progress, mock, sparse_mock_percent)
    451         concat_keys = []
    452         for tensor in tensors:
--> 453             tensor.tiles()
    454             chunk_keys = [c.key for c in tensor.chunks]
    455             result_keys.extend(chunk_keys)

~/Workspace/mars/mars/tensor/core.py in tiles(self)
    223 
    224     def tiles(self):
--> 225         return handler.tiles(self)
    226 
    227     def single_tiles(self):

~/Workspace/mars/mars/tiles.py in tiles(self, tiles_obj)
    172             if not preds or accessible:
    173                 if node.is_coarse() and node.op:
--> 174                     tiled = self._dispatch(node.op)
    175                     self._assign_to([t.data for t in tiled], node.op.outputs)
    176                 visited.add(node)

~/Workspace/mars/mars/utils.py in _wrapped(*args, **kwargs)
    352     def _wrapped(*args, **kwargs):
    353         _kernel_mode.eager = False
--> 354         return_value = func(*args, **kwargs)
    355         _kernel_mode.eager = None
    356         return return_value

~/Workspace/mars/mars/tiles.py in _dispatch(self, op)
    114             if hasattr(op_cls, 'tile'):
    115                 # has tile implementation
--> 116                 return op_cls.tile(op)
    117             for op_clz in self._handlers.keys():
    118                 if issubclass(op_cls, op_clz):

~/Workspace/mars/mars/tensor/expressions/arithmetic/core.py in tile(cls, op)
     43         for out_index in itertools.product(*(map(range, out_chunk_shape))):
     44             in_chunks = [t.cix[get_index(out_index[-t.ndim:], t)] if t.ndim != 0 else t.chunks[0]
---> 45                          for t in inputs]
     46             chunk_op = op.copy().reset_key()
     47             chunk_shape = broadcast_shape(*(c.shape for c in in_chunks))

~/Workspace/mars/mars/tensor/expressions/arithmetic/core.py in <listcomp>(.0)
     43         for out_index in itertools.product(*(map(range, out_chunk_shape))):
     44             in_chunks = [t.cix[get_index(out_index[-t.ndim:], t)] if t.ndim != 0 else t.chunks[0]
---> 45                          for t in inputs]
     46             chunk_op = op.copy().reset_key()
     47             chunk_shape = broadcast_shape(*(c.shape for c in in_chunks))

~/Workspace/mars/mars/core.py in __getitem__(self, item)
    457                 if len(item) != self._tilesable.ndim:
    458                     raise ValueError('Cannot get tensor chunk by %s, expect length %d' % (
--> 459                         item, self._tilesable.ndim))
    460 
    461                 s = self._tilesable.chunk_shape

ValueError: Cannot get tensor chunk by (0, 0), expect length 3

In [7]: a[:, mt.newaxis, :] - a[mt.newaxis, ...]                                    
Out[7]: Tensor <op=TensorSubtract, shape=(2, 2, 3), key=fe6279614688ba864daa99368caf12af>
```




Discussion:

Resolving patch:
diff --git a/mars/utils.py b/mars/utils.py
--- a/mars/utils.py
+++ b/mars/utils.py
@@ -350,9 +350,10 @@ def kernel_mode(func):
     """
 
     def _wrapped(*args, **kwargs):
-        _kernel_mode.eager = False
-        return_value = func(*args, **kwargs)
-        _kernel_mode.eager = None
-        return return_value
+        try:
+            _kernel_mode.eager = False
+            return func(*args, **kwargs)
+        finally:
+            _kernel_mode.eager = None
 
     return _wrapped


Test patch:
diff --git a/mars/tests/test_eager_mode.py b/mars/tests/test_eager_mode.py
--- a/mars/tests/test_eager_mode.py
+++ b/mars/tests/test_eager_mode.py
@@ -172,3 +172,23 @@ def testRepr(self):
 
         self.assertNotIn(repr(np.ones((10, 10))), repr(a))
         self.assertNotIn(str(np.ones((10, 10))), str(a))
+
+    def testRuntimeError(self):
+        from mars.utils import kernel_mode
+
+        @kernel_mode
+        def raise_error(*_):
+            raise ValueError
+
+        with option_context({'eager_mode': True}):
+            a = mt.zeros((10, 10))
+            with self.assertRaises(ValueError):
+                raise_error(a)
+
+            r = a + 1
+            self.assertIn(repr(np.zeros((10, 10)) + 1), repr(r))
+            np.testing.assert_array_equal(r.fetch(), np.zeros((10, 10)) + 1)
+
+        a = mt.zeros((10, 10))
+        with self.assertRaises(ValueError):
+            a.fetch()