[Placement Group] Capture Child Task Part 1 (#10968)

* In progress.

* In progers.

* Done.

* Addressed code review.

* Increase timeout to make a test less flaky.

* Addressed code review.

* Addressed code review.

python/ray/util/placement_group.py

@@ -1,4 +1,6 @@
-from typing import (List, Dict)
+import time
+
+from typing import (List, Dict, Optional)
 
 import ray
 from ray._raylet import PlacementGroupID, ObjectRef
@@ -9,11 +11,11 @@ class PlacementGroup:
 
     @staticmethod
     def empty():
-        return PlacementGroup(PlacementGroupID.nil(), [])
+        return PlacementGroup(PlacementGroupID.nil())
 
-    def __init__(self, id: PlacementGroupID, bundles: List[Dict[str, float]]):
+    def __init__(self, id: PlacementGroupID):
         self.id = id
-        self.bundles = bundles
+        self.bundle_cache = None
 
     def ready(self) -> ObjectRef:
         """Returns an ObjectRef to check ready status.
@@ -29,22 +31,22 @@ class PlacementGroup:
         >>> pg = placement_group([{"CPU": 1}])
             ray.wait([pg.ready()], timeout=0)
         """
-        worker = ray.worker.global_worker
-        worker.check_connected()
+        self._fill_bundle_cache_if_needed()
 
         @ray.remote(num_cpus=0, max_calls=0)
         def bundle_reservation_check(placement_group):
             return placement_group
 
-        assert len(self.bundles) != 0, (
+        assert len(self.bundle_cache) != 0, (
             "ready() cannot be called on placement group object with a "
-            f"bundle length == 0, current bundle length: {len(self.bundles)}")
+            "bundle length == 0, current bundle length: "
+            f"{len(self.bundle_cache)}")
 
         # Select the first bundle to schedule a dummy task.
         # Since the placement group creation will be atomic, it is sufficient
         # to schedule a single task.
         bundle_index = 0
-        bundle = self.bundles[bundle_index]
+        bundle = self.bundle_cache[bundle_index]
 
         resource_name, value = self._get_none_zero_resource(bundle)
         num_cpus = 0
@@ -67,11 +69,13 @@ class PlacementGroup:
     @property
     def bundle_specs(self) -> List[Dict]:
         """List[Dict]: Return bundles belonging to this placement group."""
-        return self.bundles
+        self._fill_bundle_cache_if_needed()
+        return self.bundle_cache
 
     @property
     def bundle_count(self):
-        return len(self.bundles)
+        self._fill_bundle_cache_if_needed()
+        return len(self.bundle_cache)
 
     def _get_none_zero_resource(self, bundle: List[Dict]):
         for key, value in bundle.items():
@@ -80,6 +84,30 @@ class PlacementGroup:
                 return key, value
         assert False, "This code should be unreachable."
 
+    def _fill_bundle_cache_if_needed(self):
+        if not self.bundle_cache:
+            # Since creating placement group is async, it is
+            # possible table is not ready yet. To avoid the
+            # problem, we should keep trying with timeout.
+            TIMEOUT_SECOND = 30
+            WAIT_INTERVAL = 0.05
+            timeout_cnt = 0
+            worker = ray.worker.global_worker
+            worker.check_connected()
+
+            while timeout_cnt < int(TIMEOUT_SECOND / WAIT_INTERVAL):
+                pg_info = ray.state.state.placement_group_table(self.id)
+                if pg_info:
+                    self.bundle_cache = list(pg_info["bundles"].values())
+                    return
+                time.sleep(WAIT_INTERVAL)
+                timeout_cnt += 1
+
+            raise RuntimeError(
+                "Couldn't get the bundle information of placement group id "
+                f"{self.id} in {TIMEOUT_SECOND} seconds. It is likely "
+                "because GCS server is too busy.")
+
 
 def placement_group(bundles: List[Dict[str, float]],
                     strategy: str = "PACK",
@@ -120,7 +148,7 @@ def placement_group(bundles: List[Dict[str, float]],
     placement_group_id = worker.core_worker.create_placement_group(
         name, bundles, strategy)
 
-    return PlacementGroup(placement_group_id, bundles)
+    return PlacementGroup(placement_group_id)
 
 
 def remove_placement_group(placement_group: PlacementGroup):
@@ -149,6 +177,41 @@ def placement_group_table(placement_group: PlacementGroup) -> dict:
     return ray.state.state.placement_group_table(placement_group.id)
 
 
+def get_current_placement_group() -> Optional[PlacementGroup]:
+    """Get the current placement group which a task or actor is using.
+
+    It returns None if there's no current placement group for the worker.
+    For example, if you call this method in your driver, it returns None
+    (because drivers never belong to any placement group).
+
+    Examples:
+
+        >>> @ray.remote
+        >>> def f():
+        >>>     # This will return the placement group the task f belongs to.
+        >>>     # It means this pg will be identical to the pg created below.
+        >>>     pg = get_current_placement_group()
+        >>> pg = placement_group([{"CPU": 2}])
+        >>> f.options(placement_group=pg).remote()
+
+        >>> # New script.
+        >>> ray.init()
+        >>> # New script doesn't belong to any placement group,
+        >>> # so it returns None.
+        >>> assert get_current_placement_group() is None
+
+    Return:
+        PlacementGroup: Placement group object.
+            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
+    if pg_id.is_nil():
+        return None
+    return PlacementGroup(pg_id)
+
+
 def check_placement_group_index(placement_group: PlacementGroup,
                                 bundle_index: int):
     assert placement_group is not None