Fix progress always is 0 or 100% (#2591)

Author: fyrestone

mars/deploy/oscar/session.py

@@ -814,16 +814,20 @@ async def _run_in_background(
             while True:
                 try:
                     if not cancelled:
-                        task_result: TaskResult = await self._task_api.wait_task(
-                            task_id, timeout=0.5
-                        )
-                        if task_result is None:
-                            # not finished, set progress
-                            progress.value = await self._task_api.get_task_progress(
-                                task_id
+                        task_result: Union[TaskResult, None] = None
+                        try:
+                            task_result = await self._task_api.wait_task(
+                                task_id, timeout=0.5
                             )
-                        else:
-                            progress.value = 1.0
+                        finally:
+                            if task_result is None:
+                                # not finished, set progress
+                                progress.value = await self._task_api.get_task_progress(
+                                    task_id
+                                )
+                            else:
+                                progress.value = 1.0
+                        if task_result is not None:
                             break
                     else:
                         # wait for task to finish

mars/deploy/oscar/tests/test_local.py

@@ -32,6 +32,7 @@
 from .... import tensor as mt
 from .... import remote as mr
 from ....config import option_context
+from ....core.context import get_context
 from ....lib.aio import new_isolation
 from ....storage import StorageLevel
 from ....services.storage import StorageAPI
@@ -378,6 +379,34 @@ def test_no_default_session():
     assert get_default_async_session() is None
 
 
+@pytest.mark.asyncio
+async def test_session_progress(create_cluster):
+    session = get_default_async_session()
+    assert session.address is not None
+    assert session.session_id is not None
+
+    def f1(interval: float, count: int):
+        for idx in range(count):
+            time.sleep(interval)
+            get_context().set_progress((1 + idx) * 1.0 / count)
+
+    r = mr.spawn(f1, args=(0.5, 10))
+
+    info = await session.execute(r)
+
+    for _ in range(20):
+        if 0 < info.progress() < 1:
+            break
+        await asyncio.sleep(0.1)
+    else:
+        raise Exception("progress test failed.")
+
+    await info
+    assert info.result() is None
+    assert info.exception() is None
+    assert info.progress() == 1
+
+
 @pytest.fixture
 def setup_session():
     session = new_session(n_cpu=2, use_uvloop=False)

mars/services/task/api/web.py

@@ -240,7 +240,9 @@ async def get_last_idle_time(self) -> Union[float, None]:
 
     async def wait_task(self, task_id: str, timeout: float = None):
         path = f"{self._address}/api/session/{self._session_id}/task/{task_id}"
-        params = {"action": "wait", "timeout": str(timeout or "")}
+        # client timeout should be longer than server timeout.
+        server_timeout = "" if timeout is None else str(max(timeout / 2.0, timeout - 1))
+        params = {"action": "wait", "timeout": server_timeout}
         res = await self._request_url(
             "GET", path, params=params, request_timeout=timeout or 0
         )