[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/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