[Ray] Implement ray task executor progress (#3008)

Author: zhongchun

mars/deploy/oscar/tests/test_ray_dag.py

@@ -116,3 +116,9 @@ async def test_iterative_tiling(ray_start_regular_shared2, create_cluster):
 @pytest.mark.parametrize("config", [{"backend": "ray"}])
 def test_sync_execute(config):
     test_local.test_sync_execute(config)
+
+
+@require_ray
+@pytest.mark.asyncio
+async def test_session_progress(ray_start_regular_shared2, create_cluster):
+    test_local.test_session_progress(create_cluster)

mars/services/task/execution/ray/executor.py

@@ -131,6 +131,12 @@ def __init__(
         self._task_context = {}
         self._available_band_resources = None
 
+        # For progress
+        self._pre_all_stages_progress = 0.0
+        self._pre_all_stages_tile_progress = 0
+        self._cur_stage_tile_progress = 0
+        self._cur_stage_output_object_refs = []
+
     @classmethod
     async def create(
         cls,
@@ -190,7 +196,12 @@ async def execute_subtask_graph(
         logger.info("Stage %s start.", stage_id)
         context = self._task_context
         output_meta_object_refs = []
-
+        self._pre_all_stages_tile_progress = (
+            self._pre_all_stages_tile_progress + self._cur_stage_tile_progress
+        )
+        self._cur_stage_tile_progress = (
+            self._tile_context.get_all_progress() - self._pre_all_stages_tile_progress
+        )
         logger.info("Submitting %s subtasks of stage %s.", len(subtask_graph), stage_id)
         result_meta_keys = {
             chunk.key
@@ -219,6 +230,7 @@ async def execute_subtask_graph(
                 continue
             elif output_count == 1:
                 output_object_refs = [output_object_refs]
+            self._cur_stage_output_object_refs.extend(output_object_refs)
             if output_meta_keys:
                 meta_object_ref, *output_object_refs = output_object_refs
                 # TODO(fyrestone): Fetch(not get) meta object here.
@@ -249,6 +261,10 @@ async def execute_subtask_graph(
 
         logger.info("Waiting for stage %s complete.", stage_id)
         ray.wait(output_object_refs, fetch_local=False)
+        # Just use `self._cur_stage_tile_progress` as current stage progress
+        # because current stage is finished, its progress is 1.
+        self._pre_all_stages_progress += self._cur_stage_tile_progress
+        self._cur_stage_output_object_refs.clear()
         logger.info("Stage %s is complete.", stage_id)
         return chunk_to_meta
 
@@ -293,7 +309,19 @@ async def get_available_band_resources(self) -> Dict[BandType, Resource]:
 
     async def get_progress(self) -> float:
         """Get the execution progress."""
-        return 1
+        stage_progress = 0.0
+        total = len(self._cur_stage_output_object_refs)
+        if total > 0:
+            finished_objects, _ = ray.wait(
+                self._cur_stage_output_object_refs,
+                num_returns=total,
+                timeout=0.1,
+                fetch_local=False,
+            )
+            stage_progress = (
+                len(finished_objects) / total * self._cur_stage_tile_progress
+            )
+        return self._pre_all_stages_progress + stage_progress
 
     async def cancel(self):
         """Cancel execution."""