[Ray] Load subtask inputs from meta (#2976)

Author: fyrestone

mars/deploy/oscar/session.py

@@ -845,7 +845,9 @@ async def init(
             from .local import new_cluster_in_isolation
 
             return (
-                await new_cluster_in_isolation(address, timeout=timeout, **kwargs)
+                await new_cluster_in_isolation(
+                    address, timeout=timeout, backend=backend, **kwargs
+                )
             ).session
 
         if kwargs:  # pragma: no cover

mars/deploy/oscar/tests/test_local.py

@@ -396,8 +396,11 @@ async def test_web_session(create_cluster, config):
     await _run_web_session_test(web_address)
 
 
-def test_sync_execute():
-    session = new_session(n_cpu=2, web=False, use_uvloop=False)
+@pytest.mark.parametrize("config", [{"backend": "mars", "incremental_index": True}])
+def test_sync_execute(config):
+    session = new_session(
+        backend=config["backend"], n_cpu=2, web=False, use_uvloop=False
+    )
 
     # web not started
     assert session._session.client.web_address is None
@@ -421,23 +424,25 @@ def test_sync_execute():
         assert d is c
         assert abs(session.fetch(d) - raw.sum()) < 0.001
 
-        with tempfile.TemporaryDirectory() as tempdir:
-            file_path = os.path.join(tempdir, "test.csv")
-            pdf = pd.DataFrame(
-                np.random.RandomState(0).rand(100, 10),
-                columns=[f"col{i}" for i in range(10)],
-            )
-            pdf.to_csv(file_path, index=False)
-
-            df = md.read_csv(file_path, chunk_bytes=os.stat(file_path).st_size / 5)
-            result = df.sum(axis=1).execute().fetch()
-            expected = pd.read_csv(file_path).sum(axis=1)
-            pd.testing.assert_series_equal(result, expected)
-
-            df = md.read_csv(file_path, chunk_bytes=os.stat(file_path).st_size / 5)
-            result = df.head(10).execute().fetch()
-            expected = pd.read_csv(file_path).head(10)
-            pd.testing.assert_frame_equal(result, expected)
+        # TODO(fyrestone): Remove this when the Ray backend support incremental index.
+        if config["incremental_index"]:
+            with tempfile.TemporaryDirectory() as tempdir:
+                file_path = os.path.join(tempdir, "test.csv")
+                pdf = pd.DataFrame(
+                    np.random.RandomState(0).rand(100, 10),
+                    columns=[f"col{i}" for i in range(10)],
+                )
+                pdf.to_csv(file_path, index=False)
+
+                df = md.read_csv(file_path, chunk_bytes=os.stat(file_path).st_size / 5)
+                result = df.sum(axis=1).execute().fetch()
+                expected = pd.read_csv(file_path).sum(axis=1)
+                pd.testing.assert_series_equal(result, expected)
+
+                df = md.read_csv(file_path, chunk_bytes=os.stat(file_path).st_size / 5)
+                result = df.head(10).execute().fetch()
+                expected = pd.read_csv(file_path).head(10)
+                pd.testing.assert_frame_equal(result, expected)
 
     for worker_pool in session._session.client._cluster._worker_pools:
         _assert_storage_cleaned(

mars/deploy/oscar/tests/test_ray_dag.py

@@ -95,3 +95,16 @@ async def create_cluster(request):
 @pytest.mark.asyncio
 async def test_execute(ray_start_regular_shared2, create_cluster, config):
     await test_local.test_execute(create_cluster, config)
+
+
+@require_ray
+@pytest.mark.asyncio
+async def test_iterative_tiling(ray_start_regular_shared2, create_cluster):
+    await test_local.test_iterative_tiling(create_cluster)
+
+
+# TODO(fyrestone): Support incremental index in ray backend.
+@require_ray
+@pytest.mark.parametrize("config", [{"backend": "ray", "incremental_index": False}])
+def test_sync_execute(config):
+    test_local.test_sync_execute(config)

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

@@ -0,0 +1,20 @@
+# Copyright 1999-2021 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+# TODO(fyrestone): Should implement the mars.core.context.Context.
+class RayExecutionContext(dict):
+    @staticmethod
+    def new_custom_log_dir():
+        return None

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

@@ -36,9 +36,10 @@
 from ....cluster.api import ClusterAPI
 from ....lifecycle.api import LifecycleAPI
 from ....meta.api import MetaAPI
-from ....subtask import SubtaskGraph
+from ....subtask import Subtask, SubtaskGraph
 from ....subtask.utils import iter_input_data_keys, iter_output_data
 from ..api import TaskExecutor, ExecutionChunkResult, register_executor_cls
+from .context import RayExecutionContext
 
 ray = lazy_import("ray")
 logger = logging.getLogger(__name__)
@@ -55,7 +56,7 @@ def execute_subtask(
     ensure_coverage()
     subtask_chunk_graph = deserialize(*subtask_chunk_graph)
     # inputs = [i[1] for i in inputs]
-    context = dict(zip(input_keys, inputs))
+    context = RayExecutionContext(zip(input_keys, inputs))
     # optimize chunk graph.
     subtask_chunk_graph = optimize(subtask_chunk_graph)
     # from data_key to results
@@ -117,7 +118,7 @@ async def create(
         session_id: str,
         address: str,
         task,
-        tile_context,
+        tile_context: TileContext,
         **kwargs
     ) -> "TaskExecutor":
         ray_executor = ray.remote(execute_subtask)
@@ -159,14 +160,10 @@ async def execute_subtask_graph(
         result_keys = {chunk.key for chunk in chunk_graph.result_chunks}
         for subtask in subtask_graph.topological_iter():
             subtask_chunk_graph = subtask.chunk_graph
-            chunk_key_to_data_keys = get_chunk_key_to_data_keys(subtask_chunk_graph)
-            key_to_input = {
-                key: context[key]
-                for key, _ in iter_input_data_keys(
-                    subtask, subtask_chunk_graph, chunk_key_to_data_keys
-                )
-            }
-            output_keys = self._get_output_keys(subtask_chunk_graph)
+            key_to_input = await self._load_subtask_inputs(
+                stage_id, subtask, subtask_chunk_graph, context
+            )
+            output_keys = self._get_subtask_output_keys(subtask_chunk_graph)
             output_meta_keys = result_keys & output_keys
             output_count = len(output_keys) + bool(output_meta_keys)
             output_object_refs = self._ray_executor.options(
@@ -250,8 +247,44 @@ async def get_progress(self) -> float:
     async def cancel(self):
         """Cancel execution."""
 
+    async def _load_subtask_inputs(
+        self, stage_id: str, subtask: Subtask, chunk_graph: ChunkGraph, context: Dict
+    ):
+        """
+        Load a dict of input key to object ref of subtask from context.
+
+        It updates the context if the input object refs are fetched from
+        the meta service.
+        """
+        key_to_input = {}
+        key_to_get_meta = {}
+        chunk_key_to_data_keys = get_chunk_key_to_data_keys(chunk_graph)
+        for key, _ in iter_input_data_keys(
+            subtask, chunk_graph, chunk_key_to_data_keys
+        ):
+            if key in context:
+                key_to_input[key] = context[key]
+            else:
+                key_to_get_meta[key] = self._meta_api.get_chunk_meta.delay(
+                    key, fields=["object_refs"]
+                )
+        if key_to_get_meta:
+            logger.info(
+                "Fetch %s metas and update context of stage %s.",
+                len(key_to_get_meta),
+                stage_id,
+            )
+            meta_list = await self._meta_api.get_chunk_meta.batch(
+                *key_to_get_meta.values()
+            )
+            for key, meta in zip(key_to_get_meta.keys(), meta_list):
+                object_ref = meta["object_refs"][0]
+                key_to_input[key] = object_ref
+                context[key] = object_ref
+        return key_to_input
+
     @staticmethod
-    def _get_output_keys(chunk_graph):
+    def _get_subtask_output_keys(chunk_graph: ChunkGraph):
         output_keys = {}
         for chunk in chunk_graph.results:
             if isinstance(chunk.op, VirtualOperand):

mars/services/task/supervisor/processor.py

@@ -204,11 +204,12 @@ async def _process_stage_chunk_graph(
         # for all execution backends.
         try:
             key_to_bands = await meta_api.get_chunk_meta.batch(*get_meta_tasks)
-        except KeyError:
-            key_to_bands = {}
-        fetch_op_to_bands = dict(
-            (key, meta["bands"][0]) for key, meta in zip(fetch_op_keys, key_to_bands)
-        )
+            fetch_op_to_bands = dict(
+                (key, meta["bands"][0])
+                for key, meta in zip(fetch_op_keys, key_to_bands)
+            )
+        except (KeyError, IndexError):
+            fetch_op_to_bands = {}
         with Timer() as timer:
             subtask_graph = await asyncio.to_thread(
                 self._preprocessor.analyze,

mars/services/task/supervisor/tests/test_task_manager.py

@@ -549,8 +549,9 @@ async def test_numexpr(actor_pool):
     ) == [1] * len(result_tileable.chunks)
 
 
+@pytest.mark.parametrize("config", [{"incremental_index": True}])
 @pytest.mark.asyncio
-async def test_optimization(actor_pool):
+async def test_optimization(actor_pool, config):
     (
         execution_backend,
         pool,
@@ -574,7 +575,7 @@ async def test_optimization(actor_pool):
         )
         pdf.to_csv(file_path, index=False)
 
-        df = md.read_csv(file_path)
+        df = md.read_csv(file_path, incremental_index=config["incremental_index"])
         df2 = df.groupby("c").agg({"a": "sum"})
         df3 = df[["b", "a"]]
 

mars/services/task/supervisor/tests/test_task_manager_on_ray_dag.py

@@ -55,3 +55,12 @@ async def test_iterative_tiling(ray_start_regular_shared2, actor_pool):
 @pytest.mark.asyncio
 async def test_numexpr(ray_start_regular_shared2, actor_pool):
     await test_task_manager.test_numexpr(actor_pool)
+
+
+# TODO(fyrestone): Support incremental index in ray backend.
+@require_ray
+@pytest.mark.parametrize("config", [{"incremental_index": False}])
+@pytest.mark.parametrize("actor_pool", [{"backend": "ray"}], indirect=True)
+@pytest.mark.asyncio
+async def test_optimization(ray_start_regular_shared2, actor_pool, config):
+    await test_task_manager.test_optimization(actor_pool, config)