[Placement Group] Capture child tasks by default. (#11025)

* In progress.

* Finished up.

* Improve comment.

* Addressed code review.

* Fix test failure.

* Fix ci failures.

* Fix CI issues.

python/ray/_raylet.pyx

@@ -796,6 +796,10 @@ cdef class CoreWorker:
             CCoreWorkerProcess.GetCoreWorker()
             .GetCurrentPlacementGroupId().Binary())
 
+    def should_capture_child_tasks_in_placement_group(self):
+        return CCoreWorkerProcess.GetCoreWorker(
+            ).ShouldCaptureChildTasksInPlacementGroup()
+
     def set_webui_display(self, key, message):
         CCoreWorkerProcess.GetCoreWorker().SetWebuiDisplay(key, message)
 
@@ -1003,7 +1007,8 @@ cdef class CoreWorker:
                     resources,
                     int max_retries,
                     PlacementGroupID placement_group_id,
-                    int64_t placement_group_bundle_index):
+                    int64_t placement_group_bundle_index,
+                    c_bool placement_group_capture_child_tasks):
         cdef:
             unordered_map[c_string, double] c_resources
             CTaskOptions task_options
@@ -1025,7 +1030,8 @@ cdef class CoreWorker:
                 CCoreWorkerProcess.GetCoreWorker().SubmitTask(
                     ray_function, args_vector, task_options, &return_ids,
                     max_retries, c_pair[CPlacementGroupID, int64_t](
-                        c_placement_group_id, placement_group_bundle_index))
+                        c_placement_group_id, placement_group_bundle_index),
+                    placement_group_capture_child_tasks)
 
             return VectorToObjectRefs(return_ids)
 
@@ -1043,6 +1049,7 @@ cdef class CoreWorker:
                      c_bool is_asyncio,
                      PlacementGroupID placement_group_id,
                      int64_t placement_group_bundle_index,
+                     c_bool placement_group_capture_child_tasks,
                      c_string extension_data
                      ):
         cdef:
@@ -1071,7 +1078,8 @@ cdef class CoreWorker:
                         dynamic_worker_options, is_detached, name, is_asyncio,
                         c_pair[CPlacementGroupID, int64_t](
                             c_placement_group_id,
-                            placement_group_bundle_index)),
+                            placement_group_bundle_index),
+                        placement_group_capture_child_tasks),
                     extension_data,
                     &c_actor_id))
 

python/ray/actor.py

@@ -6,8 +6,8 @@ import ray.ray_constants as ray_constants
 import ray._raylet
 import ray.signature as signature
 import ray.worker
-from ray.util.placement_group import PlacementGroup, \
-    check_placement_group_index
+from ray.util.placement_group import (
+    PlacementGroup, check_placement_group_index, get_current_placement_group)
 
 from ray import ActorClassID, Language
 from ray._raylet import PythonFunctionDescriptor
@@ -417,7 +417,8 @@ class ActorClass:
                 name=None,
                 lifetime=None,
                 placement_group=None,
-                placement_group_bundle_index=-1):
+                placement_group_bundle_index=-1,
+                placement_group_capture_child_tasks=None):
         """Configures and overrides the actor instantiation parameters.
 
         The arguments are the same as those that can be passed
@@ -455,7 +456,9 @@ class ActorClass:
                     name=name,
                     lifetime=lifetime,
                     placement_group=placement_group,
-                    placement_group_bundle_index=placement_group_bundle_index)
+                    placement_group_bundle_index=placement_group_bundle_index,
+                    placement_group_capture_child_tasks=(
+                        placement_group_capture_child_tasks))
 
         return ActorOptionWrapper()
 
@@ -474,7 +477,8 @@ class ActorClass:
                 name=None,
                 lifetime=None,
                 placement_group=None,
-                placement_group_bundle_index=-1):
+                placement_group_bundle_index=-1,
+                placement_group_capture_child_tasks=None):
         """Create an actor.
 
         This method allows more flexibility than the remote method because
@@ -508,6 +512,9 @@ class ActorClass:
             placement_group_bundle_index: the index of the bundle
                 if the actor belongs to a placement group, which may be -1 to
                 specify any available bundle.
+            placement_group_capture_child_tasks: Whether or not children tasks
+                of this actor should implicitly use the same placement group
+                as its parent. It is True by default.
 
         Returns:
             A handle to the newly created actor.
@@ -565,7 +572,15 @@ class ActorClass:
         else:
             raise ValueError("lifetime must be either `None` or 'detached'")
 
+        if placement_group_capture_child_tasks is None:
+            placement_group_capture_child_tasks = (
+                worker.should_capture_child_tasks_in_placement_group)
+
         if placement_group is None:
+            if placement_group_capture_child_tasks:
+                placement_group = get_current_placement_group()
+
+        if not placement_group:
             placement_group = PlacementGroup.empty()
 
         check_placement_group_index(placement_group,
@@ -644,6 +659,7 @@ class ActorClass:
             is_asyncio,
             placement_group.id,
             placement_group_bundle_index,
+            placement_group_capture_child_tasks,
             # Store actor_method_cpu in actor handle's extension data.
             extension_data=str(actor_method_cpu))
 

python/ray/cross_language.py

@@ -78,7 +78,8 @@ def java_function(class_name, function_name):
         None,  # max_retries
         placement_group=None,
         # TODO(ekl) set default to -1 once we support -1 as "any index"
-        placement_group_bundle_index=0)
+        placement_group_bundle_index=0,
+        placement_group_capture_child_tasks=None)
 
 
 def java_actor_class(class_name):

python/ray/includes/common.pxd

@@ -254,7 +254,8 @@ cdef extern from "ray/core_worker/common.h" nogil:
             const unordered_map[c_string, double] &placement_resources,
             const c_vector[c_string] &dynamic_worker_options,
             c_bool is_detached, c_string &name, c_bool is_asyncio,
-            c_pair[CPlacementGroupID, int64_t] placement_options)
+            c_pair[CPlacementGroupID, int64_t] placement_options,
+            c_bool placement_group_capture_child_tasks)
 
     cdef cppclass CPlacementGroupCreationOptions \
             "ray::PlacementGroupCreationOptions":

python/ray/includes/libcoreworker.pxd

@@ -89,7 +89,8 @@ cdef extern from "ray/core_worker/core_worker.h" nogil:
             const c_vector[unique_ptr[CTaskArg]] &args,
             const CTaskOptions &options, c_vector[CObjectID] *return_ids,
             int max_retries,
-            c_pair[CPlacementGroupID, int64_t] placement_options)
+            c_pair[CPlacementGroupID, int64_t] placement_options,
+            c_bool placement_group_capture_child_tasks)
         CRayStatus CreateActor(
             const CRayFunction &function,
             const c_vector[unique_ptr[CTaskArg]] &args,
@@ -123,6 +124,7 @@ cdef extern from "ray/core_worker/core_worker.h" nogil:
         CTaskID GetCurrentTaskId()
         CNodeID GetCurrentNodeId()
         CPlacementGroupID GetCurrentPlacementGroupId()
+        c_bool ShouldCaptureChildTasksInPlacementGroup()
         const CActorID &GetActorId()
         void SetActorTitle(const c_string &title)
         void SetWebuiDisplay(const c_string &key, const c_string &message)

python/ray/remote_function.py

@@ -4,8 +4,11 @@ from functools import wraps
 from ray import cloudpickle as pickle
 from ray._raylet import PythonFunctionDescriptor
 from ray import cross_language, Language
-from ray.util.placement_group import PlacementGroup, \
-    check_placement_group_index
+from ray.util.placement_group import (
+    PlacementGroup,
+    check_placement_group_index,
+    get_current_placement_group,
+)
 import ray.signature
 
 # Default parameters for remote functions.
@@ -63,7 +66,8 @@ class RemoteFunction:
     def __init__(self, language, function, function_descriptor, num_cpus,
                  num_gpus, memory, object_store_memory, resources,
                  accelerator_type, num_returns, max_calls, max_retries,
-                 placement_group, placement_group_bundle_index):
+                 placement_group, placement_group_bundle_index,
+                 placement_group_capture_child_tasks):
         self._language = language
         self._function = function
         self._function_name = (
@@ -135,6 +139,7 @@ class RemoteFunction:
                 max_retries=None,
                 placement_group=None,
                 placement_group_bundle_index=-1,
+                placement_group_capture_child_tasks=None,
                 name=""):
         """Configures and overrides the task invocation parameters.
 
@@ -168,6 +173,8 @@ class RemoteFunction:
                     max_retries=max_retries,
                     placement_group=placement_group,
                     placement_group_bundle_index=placement_group_bundle_index,
+                    placement_group_capture_child_tasks=(
+                        placement_group_capture_child_tasks),
                     name=name)
 
         return FuncWrapper()
@@ -185,6 +192,7 @@ class RemoteFunction:
                 max_retries=None,
                 placement_group=None,
                 placement_group_bundle_index=-1,
+                placement_group_capture_child_tasks=None,
                 name=""):
         """Submit the remote function for execution."""
         worker = ray.worker.global_worker
@@ -220,7 +228,15 @@ class RemoteFunction:
         if max_retries is None:
             max_retries = self._max_retries
 
+        if placement_group_capture_child_tasks is None:
+            placement_group_capture_child_tasks = (
+                worker.should_capture_child_tasks_in_placement_group)
+
         if placement_group is None:
+            if placement_group_capture_child_tasks:
+                placement_group = get_current_placement_group()
+
+        if not placement_group:
             placement_group = PlacementGroup.empty()
 
         check_placement_group_index(placement_group,
@@ -248,7 +264,8 @@ class RemoteFunction:
             object_refs = worker.core_worker.submit_task(
                 self._language, self._function_descriptor, list_args, name,
                 num_returns, resources, max_retries, placement_group.id,
-                placement_group_bundle_index)
+                placement_group_bundle_index,
+                placement_group_capture_child_tasks)
 
             if len(object_refs) == 1:
                 return object_refs[0]

python/ray/runtime_context.py

@@ -54,6 +54,18 @@ class RuntimeContext(object):
         """
         return self.worker.placement_group_id
 
+    @property
+    def should_capture_child_tasks_in_placement_group(self):
+        """Get if the current task should capture parent's placement group.
+
+        This returns True if it is called inside a driver.
+
+        Returns:
+            Return True if the current task should implicitly
+                capture the parent placement group.
+        """
+        return self.worker.should_capture_child_tasks_in_placement_group
+
 
 _runtime_context = None
 

python/ray/tests/BUILD

@@ -52,7 +52,6 @@ py_test_module_list(
     "test_multiprocessing.py",
     "test_object_manager.py",
     "test_output.py",
-    "test_placement_group.py",
     "test_reconstruction.py",
     "test_reference_counting_2.py",
     "test_reference_counting.py",
@@ -104,7 +103,8 @@ py_test_module_list(
 py_test_module_list(
   files = [
     "test_stress_failure.py",
-    "test_failure.py"
+    "test_failure.py",
+    "test_placement_group.py",
   ],
   size = "large",
   extra_srcs = SRCS,

python/ray/tests/test_placement_group.py

@@ -793,7 +793,7 @@ def test_mini_integration(ray_start_cluster):
     assert all(ray.get([a.ping.remote() for a in actors]))
 
 
-def test_capture_child_tasks(ray_start_cluster):
+def test_capture_child_actors(ray_start_cluster):
     cluster = ray_start_cluster
     total_num_actors = 4
     for _ in range(2):
@@ -806,12 +806,13 @@ def test_capture_child_tasks(ray_start_cluster):
         }, {
             "CPU": 2
         }], strategy="STRICT_PACK")
-    ray.get(pg.ready(), timeout=5)
+    ray.get(pg.ready())
 
     # If get_current_placement_group is used when the current worker/driver
     # doesn't belong to any of placement group, it should return None.
     assert get_current_placement_group() is None
 
+    # Test actors first.
     @ray.remote(num_cpus=1)
     class NestedActor:
         def ready(self):
@@ -826,8 +827,16 @@ def test_capture_child_tasks(ray_start_cluster):
             return True
 
         def schedule_nested_actor(self):
-            actor = NestedActor.options(
-                placement_group=get_current_placement_group()).remote()
+            # Make sure we can capture the current placement group.
+            assert get_current_placement_group() is not None
+            # Actors should be implicitly captured.
+            actor = NestedActor.remote()
+            ray.get(actor.ready.remote())
+            self.actors.append(actor)
+
+        def schedule_nested_actor_outside_pg(self):
+            # Don't use placement group.
+            actor = NestedActor.options(placement_group=None).remote()
             ray.get(actor.ready.remote())
             self.actors.append(actor)
 
@@ -846,6 +855,104 @@ def test_capture_child_tasks(ray_start_cluster):
     # Since all node id should be identical, set should be equal to 1.
     assert len(node_id_set) == 1
 
+    # Kill an actor and wait until it is killed.
+    ray.kill(a)
+    with pytest.raises(ray.exceptions.RayActorError):
+        ray.get(a.ready.remote())
+
+    # Now create an actor, but do not capture the current tasks
+    a = Actor.options(
+        placement_group=pg,
+        placement_group_capture_child_tasks=False).remote()
+    ray.get(a.ready.remote())
+    # 1 top level actor + 3 children.
+    for _ in range(total_num_actors - 1):
+        ray.get(a.schedule_nested_actor.remote())
+    # Make sure all the actors are not scheduled on the same node.
+    # It is because the child tasks are not scheduled on the same
+    # placement group.
+    node_id_set = set()
+    for actor_info in ray.actors().values():
+        node_id = actor_info["Address"]["NodeID"]
+        node_id_set.add(node_id)
+
+    assert len(node_id_set) == 2
+
+    # Kill an actor and wait until it is killed.
+    ray.kill(a)
+    with pytest.raises(ray.exceptions.RayActorError):
+        ray.get(a.ready.remote())
+
+    # Lastly, make sure when None is specified, actors are not scheduled
+    # on the same placement group.
+    a = Actor.options(placement_group=pg).remote()
+    ray.get(a.ready.remote())
+    # 1 top level actor + 3 children.
+    for _ in range(total_num_actors - 1):
+        ray.get(a.schedule_nested_actor_outside_pg.remote())
+    # Make sure all the actors are not scheduled on the same node.
+    # It is because the child tasks are not scheduled on the same
+    # placement group.
+    node_id_set = set()
+    for actor_info in ray.actors().values():
+        node_id = actor_info["Address"]["NodeID"]
+        node_id_set.add(node_id)
+
+    assert len(node_id_set) == 2
+
+
+def test_capture_child_tasks(ray_start_cluster):
+    cluster = ray_start_cluster
+    total_num_tasks = 4
+    for _ in range(2):
+        cluster.add_node(num_cpus=total_num_tasks, num_gpus=total_num_tasks)
+    ray.init(address=cluster.address)
+
+    pg = ray.util.placement_group(
+        [{
+            "CPU": 2,
+            "GPU": 2,
+        }, {
+            "CPU": 2,
+            "GPU": 2,
+        }],
+        strategy="STRICT_PACK")
+    ray.get(pg.ready())
+
+    # If get_current_placement_group is used when the current worker/driver
+    # doesn't belong to any of placement group, it should return None.
+    assert get_current_placement_group() is None
+
+    # Test if tasks capture child tasks.
+    @ray.remote
+    def task():
+        return get_current_placement_group()
+
+    @ray.remote
+    def create_nested_task(child_cpu, child_gpu):
+        assert get_current_placement_group() is not None
+        return ray.get([
+            task.options(num_cpus=child_cpu, num_gpus=child_gpu).remote()
+            for _ in range(3)
+        ])
+
+    t = create_nested_task.options(
+        num_cpus=1, num_gpus=0, placement_group=pg).remote(1, 0)
+    pgs = ray.get(t)
+    # Every task should have current placement group because they
+    # should be implicitly captured by default.
+    assert None not in pgs
+
+    # Test if tasks don't capture child tasks when the option is off.
+    t2 = create_nested_task.options(
+        num_cpus=0,
+        num_gpus=1,
+        placement_group=pg,
+        placement_group_capture_child_tasks=False).remote(0, 1)
+    pgs = ray.get(t2)
+    # All placement group should be None because we don't capture child tasks.
+    assert not all(pgs)
+
 
 def test_ready_warning_suppressed(ray_start_regular, error_pubsub):
     p = error_pubsub

python/ray/util/placement_group.py

@@ -209,8 +209,9 @@ def get_current_placement_group() -> Optional[PlacementGroup]:
             None if the current task or actor wasn't
             created with any placement group.
     """
-    pg_id = ray.runtime_context.get_runtime_context(
-    ).current_placement_group_id
+    worker = ray.worker.global_worker
+    worker.check_connected()
+    pg_id = worker.placement_group_id
     if pg_id.is_nil():
         return None
     return PlacementGroup(pg_id)

python/ray/worker.py

@@ -159,6 +159,10 @@ class Worker:
     def placement_group_id(self):
         return self.core_worker.get_placement_group_id()
 
+    @property
+    def should_capture_child_tasks_in_placement_group(self):
+        return self.core_worker.should_capture_child_tasks_in_placement_group()
+
     @property
     def current_session_and_job(self):
         """Get the current session index and job id as pair."""
@@ -1671,7 +1675,8 @@ def make_decorator(num_returns=None,
                    max_task_retries=None,
                    worker=None,
                    placement_group=None,
-                   placement_group_bundle_index=-1):
+                   placement_group_bundle_index=-1,
+                   placement_group_capture_child_tasks=True):
     def decorator(function_or_class):
         if (inspect.isfunction(function_or_class)
                 or is_cython(function_or_class)):
@@ -1701,7 +1706,8 @@ def make_decorator(num_returns=None,
                 Language.PYTHON, function_or_class, None, num_cpus, num_gpus,
                 memory, object_store_memory, resources, accelerator_type,
                 num_returns, max_calls, max_retries, placement_group,
-                placement_group_bundle_index)
+                placement_group_bundle_index,
+                placement_group_capture_child_tasks)
 
         if inspect.isclass(function_or_class):
             if num_returns is not None:
@@ -1831,6 +1837,9 @@ def remote(*args, **kwargs):
         placement_group_bundle_index (int): The index of the bundle
             if the task belongs to a placement group, which may be
             -1 to indicate any available bundle.
+        placement_group_capture_child_tasks (bool): Default True.
+            If True, all the child tasks (including actor creation)
+            are scheduled in the same placement group.
 
     """
     worker = global_worker
@@ -1864,6 +1873,7 @@ def remote(*args, **kwargs):
             "max_retries",
             "placement_group",
             "placement_group_bundle_index",
+            "placement_group_capture_child_tasks",
         ], error_string
 
     num_cpus = kwargs["num_cpus"] if "num_cpus" in kwargs else None