[Ray] Create RayTaskState actor as needed by default (#3081)

Author: fyrestone

.github/workflows/platform-ci.yml

@@ -94,7 +94,7 @@ jobs:
               rm -fr /tmp/etcd-$ETCD_VER-linux-amd64.tar.gz /tmp/etcd-download-test
             fi
             if [ -n "$WITH_RAY" ] || [ -n "$WITH_RAY_DAG" ]; then
-              pip install ray[default]==1.9.2
+              pip install ray[default]==1.9.2 "protobuf<4"
               pip install "xgboost_ray==0.1.5" "xgboost<1.6.0"
             fi
             if [ -n "$RUN_DASK" ]; then

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

@@ -47,3 +47,11 @@ def get_n_worker(self):
 
     def get_subtask_cancel_timeout(self):
         return self._ray_execution_config.get("subtask_cancel_timeout")
+
+    def create_task_state_actor_as_needed(self):
+        # Whether create RayTaskState actor as needed.
+        #   - True (default):
+        #     Create RayTaskState actor only when create_remote_object is called.
+        #   - False:
+        #     Create RayTaskState actor in advance when the RayTaskExecutor is created.
+        return self._ray_execution_config.get("create_task_state_actor_as_needed", True)

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

@@ -16,7 +16,7 @@
 import inspect
 import logging
 from dataclasses import asdict
-from typing import Union, Dict, List
+from typing import Dict, List, Callable
 
 from .....core.context import Context
 from .....storage.base import StorageLevel
@@ -77,18 +77,21 @@ def wrap(*args, **kwargs):
 
 class _RayRemoteObjectContext:
     def __init__(
-        self, actor_name_or_handle: Union[str, "ray.actor.ActorHandle"], *args, **kwargs
+        self,
+        get_or_create_actor: Callable[[], "ray.actor.ActorHandle"],
+        *args,
+        **kwargs
     ):
         super().__init__(*args, **kwargs)
-        self._actor_name_or_handle = actor_name_or_handle
+        self._get_or_create_actor = get_or_create_actor
         self._task_state_actor = None
 
     def _get_task_state_actor(self) -> "ray.actor.ActorHandle":
+        # Get the RayTaskState actor, this is more clear and faster than wraps
+        # the `get_or_create_actor` by lru_cache in __init__ because this method
+        # is called as needed.
         if self._task_state_actor is None:
-            if isinstance(self._actor_name_or_handle, ray.actor.ActorHandle):
-                self._task_state_actor = self._actor_name_or_handle
-            else:
-                self._task_state_actor = ray.get_actor(self._actor_name_or_handle)
+            self._task_state_actor = self._get_or_create_actor()
         return self._task_state_actor
 
     @implements(Context.create_remote_object)
@@ -124,13 +127,15 @@ def __init__(
         config: RayExecutionConfig,
         task_context: Dict,
         task_chunks_meta: Dict,
+        worker_addresses: List[str],
         *args,
         **kwargs
     ):
         super().__init__(*args, **kwargs)
         self._config = config
         self._task_context = task_context
         self._task_chunks_meta = task_chunks_meta
+        self._worker_addresses = worker_addresses
 
     @implements(Context.get_chunks_result)
     def get_chunks_result(self, data_keys: List[str]) -> List:
@@ -158,6 +163,11 @@ def get_total_n_cpu(self) -> int:
         # TODO(fyrestone): Support auto scaling.
         return self._config.get_n_cpu() * self._config.get_n_worker()
 
+    @implements(Context.get_worker_addresses)
+    def get_worker_addresses(self) -> List[str]:
+        # Returns virtual worker addresses.
+        return self._worker_addresses
+
 
 # TODO(fyrestone): Implement more APIs for Ray.
 class RayExecutionWorkerContext(_RayRemoteObjectContext, dict):

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

@@ -15,8 +15,9 @@
 import asyncio
 import functools
 import logging
+import operator
 from dataclasses import dataclass
-from typing import List, Dict, Any, Set
+from typing import List, Dict, Any, Set, Callable
 from .....core import ChunkGraph, Chunk, TileContext
 from .....core.context import set_context
 from .....core.operand import (
@@ -65,9 +66,23 @@ class _RayChunkMeta:
 
 class RayTaskState(RayRemoteObjectManager):
     @classmethod
-    def gen_name(cls, task_id: str):
+    def _gen_name(cls, task_id: str):
         return f"{cls.__name__}_{task_id}"
 
+    @classmethod
+    def get_handle(cls, task_id: str):
+        """Get the RayTaskState actor handle."""
+        name = cls._gen_name(task_id)
+        logger.info("Getting %s handle.", name)
+        return ray.get_actor(name)
+
+    @classmethod
+    def create(cls, task_id: str):
+        """Create a RayTaskState actor."""
+        name = cls._gen_name(task_id)
+        logger.info("Creating %s.", name)
+        return ray.remote(cls).options(name=name).remote()
+
 
 _optimize_physical = None
 
@@ -109,7 +124,7 @@ def execute_subtask(
     subtask_chunk_graph = deserialize(*subtask_chunk_graph)
     # inputs = [i[1] for i in inputs]
     context = RayExecutionWorkerContext(
-        RayTaskState.gen_name(task_id), zip(input_keys, inputs)
+        lambda: RayTaskState.get_handle(task_id), zip(input_keys, inputs)
     )
     # optimize chunk graph.
     subtask_chunk_graph = _optimize_subtask_graph(subtask_chunk_graph)
@@ -152,7 +167,6 @@ def __init__(
         tile_context: TileContext,
         task_context: Dict[str, "ray.ObjectRef"],
         task_chunks_meta: Dict[str, _RayChunkMeta],
-        task_state_actor: "ray.actor.ActorHandle",
         lifecycle_api: LifecycleAPI,
         meta_api: MetaAPI,
     ):
@@ -161,7 +175,6 @@ def __init__(
         self._tile_context = tile_context
         self._task_context = task_context
         self._task_chunks_meta = task_chunks_meta
-        self._task_state_actor = task_state_actor
         self._ray_executor = self._get_ray_executor()
 
         # api
@@ -196,31 +209,39 @@ async def create(
         **kwargs,
     ) -> "RayTaskExecutor":
         lifecycle_api, meta_api = await cls._get_apis(session_id, address)
-        task_state_actor = (
-            ray.remote(RayTaskState)
-            .options(name=RayTaskState.gen_name(task.task_id))
-            .remote()
-        )
         task_context = {}
         task_chunks_meta = {}
-        await cls._init_context(
-            config,
-            task_context,
-            task_chunks_meta,
-            task_state_actor,
-            session_id,
-            address,
-        )
-        return cls(
+
+        executor = cls(
             config,
             task,
             tile_context,
             task_context,
             task_chunks_meta,
-            task_state_actor,
             lifecycle_api,
             meta_api,
         )
+        available_band_resources = await executor.get_available_band_resources()
+        worker_addresses = list(
+            map(operator.itemgetter(0), available_band_resources.keys())
+        )
+        if config.create_task_state_actor_as_needed():
+            create_task_state_actor = lambda: RayTaskState.create(  # noqa: E731
+                task_id=task.task_id
+            )
+        else:
+            actor_handle = RayTaskState.create(task_id=task.task_id)
+            create_task_state_actor = lambda: actor_handle  # noqa: E731
+        await cls._init_context(
+            config,
+            task_context,
+            task_chunks_meta,
+            create_task_state_actor,
+            worker_addresses,
+            session_id,
+            address,
+        )
+        return executor
 
     # noinspection DuplicatedCode
     def destroy(self):
@@ -229,7 +250,6 @@ def destroy(self):
         self._tile_context = None
         self._task_context = None
         self._task_chunks_meta = None
-        self._task_state_actor = None
         self._ray_executor = None
 
         # api
@@ -267,7 +287,8 @@ async def _init_context(
         config: RayExecutionConfig,
         task_context: Dict[str, "ray.ObjectRef"],
         task_chunks_meta: Dict[str, _RayChunkMeta],
-        task_state_actor: "ray.actor.ActorHandle",
+        create_task_state_actor: Callable[[], "ray.actor.ActorHandle"],
+        worker_addresses: List[str],
         session_id: str,
         address: str,
     ):
@@ -276,7 +297,8 @@ async def _init_context(
             config,
             task_context,
             task_chunks_meta,
-            task_state_actor,
+            worker_addresses,
+            create_task_state_actor,
             session_id,
             address,
             address,
@@ -426,7 +448,9 @@ async def get_available_band_resources(self) -> Dict[BandType, Resource]:
             idx = 0
             for band_resource in band_resources:
                 for band, resource in band_resource.items():
-                    virtual_band_resources[(f"ray_virtual://{idx}", band)] = resource
+                    virtual_band_resources[
+                        (f"ray_virtual_address_{idx}:0", band)
+                    ] = resource
                     idx += 1
             self._available_band_resources = virtual_band_resources
 

mars/services/task/execution/ray/tests/test_ray_execution_backend.py

@@ -21,6 +21,7 @@
 from ...... import tensor as mt
 
 from ......core import TileContext
+from ......core.context import get_context
 from ......core.graph import TileableGraph, TileableGraphBuilder, ChunkGraphBuilder
 from ......serialization import serialize
 from ......tests.core import require_ray, mock
@@ -34,7 +35,7 @@
     RayRemoteObjectManager,
     _RayRemoteObjectContext,
 )
-from ..executor import execute_subtask, RayTaskExecutor
+from ..executor import execute_subtask, RayTaskExecutor, RayTaskState
 from ..fetcher import RayFetcher
 
 ray = lazy_import("ray")
@@ -51,11 +52,18 @@ def __init__(self, *args, **kwargs):
         self._set_attrs = Counter()
         super().__init__(*args, **kwargs)
 
+    @classmethod
+    async def _get_apis(cls, session_id: str, address: str):
+        return None, None
+
     @staticmethod
     def _get_ray_executor():
         # Export remote function once.
         return None
 
+    async def get_available_band_resources(self):
+        return {}
+
     def set_attr_counter(self):
         return self._set_attrs
 
@@ -64,6 +72,54 @@ def __setattr__(self, key, value):
         self._set_attrs[key] += 1
 
 
+@require_ray
+@pytest.mark.asyncio
+@mock.patch("mars.services.task.execution.ray.executor.RayTaskState.create")
+@mock.patch("mars.services.task.execution.ray.context.RayExecutionContext.init")
+@mock.patch("ray.get")
+async def test_ray_executor_create(
+    mock_ray_get, mock_execution_context_init, mock_task_state_actor_create
+):
+    task = Task("mock_task", "mock_session")
+
+    # Create RayTaskState actor as needed by default.
+    mock_config = RayExecutionConfig.from_execution_config({"backend": "ray"})
+    executor = await MockRayTaskExecutor.create(
+        mock_config,
+        session_id="mock_session_id",
+        address="mock_address",
+        task=task,
+        tile_context=TileContext(),
+    )
+    assert isinstance(executor, MockRayTaskExecutor)
+    assert mock_task_state_actor_create.call_count == 0
+    ctx = get_context()
+    assert isinstance(ctx, RayExecutionContext)
+    ctx.create_remote_object("abc", lambda: None)
+    assert mock_ray_get.call_count == 1
+    assert mock_task_state_actor_create.call_count == 1
+
+    # Create RayTaskState actor in advance if create_task_state_actor_as_needed is False
+    mock_config = RayExecutionConfig.from_execution_config(
+        {"backend": "ray", "ray": {"create_task_state_actor_as_needed": False}}
+    )
+    executor = await MockRayTaskExecutor.create(
+        mock_config,
+        session_id="mock_session_id",
+        address="mock_address",
+        task=task,
+        tile_context=TileContext(),
+    )
+    assert isinstance(executor, MockRayTaskExecutor)
+    assert mock_ray_get.call_count == 1
+    assert mock_task_state_actor_create.call_count == 2
+    ctx = get_context()
+    assert isinstance(ctx, RayExecutionContext)
+    ctx.create_remote_object("abc", lambda: None)
+    assert mock_ray_get.call_count == 2
+    assert mock_task_state_actor_create.call_count == 2
+
+
 def test_ray_executor_destroy():
     task = Task("mock_task", "mock_session")
     mock_config = RayExecutionConfig.from_execution_config({"backend": "ray"})
@@ -73,7 +129,6 @@ def test_ray_executor_destroy():
         tile_context=TileContext(),
         task_context={},
         task_chunks_meta={},
-        task_state_actor=None,
         lifecycle_api=None,
         meta_api=None,
     )
@@ -163,8 +218,8 @@ async def bar(self, a, b):
         await manager.call_remote_object(name, "foo", 3, 4)
 
     # Test _RayRemoteObjectContext
-    remote_manager = ray.remote(RayRemoteObjectManager).remote()
-    context = _RayRemoteObjectContext(remote_manager)
+    test_task_id = "test_task_id"
+    context = _RayRemoteObjectContext(lambda: RayTaskState.create(test_task_id))
     context.create_remote_object(name, _TestRemoteObject, 2)
     remote_object = context.get_remote_object(name)
     r = remote_object.foo(3, 4)
@@ -185,9 +240,12 @@ def __init__(self):
     with pytest.raises(MyException):
         context.create_remote_object(name, _ErrorRemoteObject)
 
+    handle = RayTaskState.get_handle(test_task_id)
+    assert handle is not None
+
 
 @require_ray
-def test_get_chunks_result(ray_start_regular_shared2):
+def test_ray_execution_context(ray_start_regular_shared2):
     value = 123
     o = ray.put(value)
 
@@ -196,13 +254,19 @@ def fake_init(self):
 
     with mock.patch.object(ThreadedServiceContext, "__init__", new=fake_init):
         mock_config = RayExecutionConfig.from_execution_config({"backend": "ray"})
-        context = RayExecutionContext(mock_config, {"abc": o}, {}, None)
+        mock_worker_addresses = ["mock_worker_address"]
+        context = RayExecutionContext(
+            mock_config, {"abc": o}, {}, mock_worker_addresses, lambda: None
+        )
         r = context.get_chunks_result(["abc"])
         assert r == [value]
 
+        r = context.get_worker_addresses()
+        assert r == mock_worker_addresses
+
 
 def test_ray_execution_worker_context():
-    context = RayExecutionWorkerContext(None)
+    context = RayExecutionWorkerContext(lambda: None)
     with pytest.raises(NotImplementedError):
         context.set_running_operand_key("mock_session_id", "mock_op_key")
     with pytest.raises(NotImplementedError):