[BACKPORT]Fix duplicate exceptions in log (#2723) (#2736)

Co-authored-by: Liu Bao <[email protected]>

Author: wjsi

mars/deploy/oscar/session.py

@@ -124,6 +124,10 @@ def __await__(self):
         self._ensure_future()
         return self._future_local.aio_future.__await__()
 
+    def get_future(self):
+        self._ensure_future()
+        return self._future_local.aio_future
+
 
 warning_msg = """
 No session found, local session \
@@ -1670,6 +1674,27 @@ def __exit__(self, *_):
         self.close()
 
 
+async def _execute_with_progress(
+    execution_info: ExecutionInfo,
+    progress_bar: ProgressBar,
+    progress_update_interval: Union[int, float],
+    cancelled: asyncio.Event,
+):
+    with progress_bar:
+        while not cancelled.is_set():
+            done, _pending = await asyncio.wait(
+                [execution_info.get_future()], timeout=progress_update_interval
+            )
+            if not done:
+                if not cancelled.is_set() and execution_info.progress() is not None:
+                    progress_bar.update(execution_info.progress() * 100)
+            else:
+                # done
+                if not cancelled.is_set():
+                    progress_bar.update(100)
+                break
+
+
 async def _execute(
     *tileables: Tuple[TileableType],
     session: _IsolatedSession = None,
@@ -1691,39 +1716,20 @@ def _attach_session(future: asyncio.Future):
     if wait:
         progress_bar = ProgressBar(show_progress)
         if progress_bar.show_progress:
-            with progress_bar:
-                while not cancelled.is_set():
-                    try:
-                        await asyncio.wait_for(
-                            asyncio.shield(execution_info), progress_update_interval
-                        )
-                        # done
-                        if not cancelled.is_set():
-                            progress_bar.update(100)
-                        break
-                    except asyncio.TimeoutError:
-                        # timeout
-                        if (
-                            not cancelled.is_set()
-                            and execution_info.progress() is not None
-                        ):
-                            progress_bar.update(execution_info.progress() * 100)
-                if cancelled.is_set():
-                    # cancel execution
-                    execution_info.cancel()
-                    execution_info.remove_done_callback(_attach_session)
-                    await execution_info
+            await _execute_with_progress(
+                execution_info, progress_bar, progress_update_interval, cancelled
+            )
         else:
             await asyncio.wait(
                 [execution_info, cancelled.wait()], return_when=asyncio.FIRST_COMPLETED
             )
-            if cancelled.is_set():
-                execution_info.remove_done_callback(_attach_session)
-                execution_info.cancel()
-            else:
-                # set cancelled to avoid wait task leak
-                cancelled.set()
-            await execution_info
+        if cancelled.is_set():
+            execution_info.remove_done_callback(_attach_session)
+            execution_info.cancel()
+        else:
+            # set cancelled to avoid wait task leak
+            cancelled.set()
+        await execution_info
     else:
         return execution_info
 

mars/deploy/oscar/tests/test_local.py

@@ -37,6 +37,7 @@
 from ....storage import StorageLevel
 from ....services.storage import StorageAPI
 from ....tensor.arithmetic.add import TensorAdd
+from ....tests.core import mock
 from ..local import new_cluster
 from ..service import load_config
 from ..session import (
@@ -48,7 +49,10 @@
     fetch_infos,
     stop_server,
     AsyncSession,
+    ExecutionInfo,
+    Progress,
     _IsolatedWebSession,
+    _execute_with_progress,
 )
 from .modules.utils import (  # noqa: F401; pylint: disable=unused-variable
     cleanup_third_party_modules_output,
@@ -574,3 +578,36 @@ def test_load_third_party_modules(cleanup_third_party_modules_output):  # noqa:
 
     session.stop_server()
     assert get_default_session() is None
+
+
+@mock.patch("asyncio.base_events.logger")
+def test_show_progress_raise_exception(m_log):
+    loop = asyncio.get_event_loop()
+    event = asyncio.Event()
+
+    class ProgressBar:
+        def __init__(self, *args, **kwargs):
+            pass
+
+        def __enter__(self):
+            pass
+
+        def __exit__(self, *_):
+            pass
+
+        def update(self, progress: float):
+            pass
+
+    async def _exec():
+        progress = Progress()
+        execution_info = ExecutionInfo(
+            asyncio.create_task(event.wait()), progress, loop
+        )
+        progress_bar = ProgressBar(True)
+        cancel_event = asyncio.Event()
+        loop.call_later(2, cancel_event.set)
+        await _execute_with_progress(execution_info, progress_bar, 0.01, cancel_event)
+        execution_info.get_future().set_exception(Exception("Expect Exception!!!"))
+
+    loop.run_until_complete(_exec())
+    assert len(m_log.mock_calls) < 3