Persist GangContext in ReplicaContext

Signed-off-by: jeffreywang <jeffreywang@anyscale.com>

Author: jeffreywang-anyscale

python/ray/serve/_private/common.py

@@ -859,6 +859,24 @@ class CreatePlacementGroupRequest:
     fallback_strategy: Optional[List[Dict[str, Any]]] = None
 
 
+@PublicAPI(stability="alpha")
+@dataclass
+class GangContext:
+    """Context information for a replica that is part of a gang.
+
+    Attributes:
+        gang_id: Unique identifier for this gang.
+        rank: This replica's rank within the gang (0-indexed).
+        world_size: Total number of replicas in this gang.
+        member_replica_ids: List of replica IDs in this gang, ordered by rank.
+    """
+
+    gang_id: str
+    rank: int
+    world_size: int
+    member_replica_ids: List[str]
+
+
 @dataclass
 class GangPlacementGroupRequest:
     """Request to prepare gang placement groups for a deployment.

python/ray/serve/_private/deployment_state.py

@@ -28,6 +28,7 @@
     DeploymentStatusTrigger,
     DeploymentTargetInfo,
     Duration,
+    GangContext,
     GangPlacementGroupRequest,
     GangPreparationResult,
     ReplicaID,
@@ -400,6 +401,10 @@ def deployment_name(self) -> str:
     def rank(self) -> Optional[ReplicaRank]:
         return self._rank
 
+    @property
+    def gang_context(self) -> Optional[GangContext]:
+        return self._gang_context
+
     @property
     def app_name(self) -> str:
         return self._deployment_id.app_name
@@ -578,6 +583,7 @@ def start(
         assign_rank_callback: Callable[[ReplicaID], ReplicaRank],
         gang_placement_group: Optional[PlacementGroup] = None,
         gang_bundle_index: Optional[int] = None,
+        gang_context: Optional[GangContext] = None,
     ) -> ReplicaSchedulingRequest:
         """Start the current DeploymentReplica instance.
 
@@ -589,12 +595,14 @@ def start(
             assign_rank_callback: Callback to assign rank to the replica.
             gang_placement_group: Pre-created gang PG to schedule this replica on.
             gang_bundle_index: Bundle index within the gang PG for this replica.
+            gang_context: Gang context for this replica (if part of a gang).
         """
         self._assign_rank_callback = assign_rank_callback
         self._actor_resources = deployment_info.replica_config.resource_dict
         self._ingress = deployment_info.ingress
         self._gang_placement_group = gang_placement_group
         self._gang_bundle_index = gang_bundle_index
+        self._gang_context = gang_context
         # it is currently not possible to create a placement group
         # with no resources (https://github.com/ray-project/ray/issues/20401)
         self._deployment_is_cross_language = (
@@ -894,7 +902,7 @@ def check_ready(self) -> Tuple[ReplicaStartupStatus, Optional[str]]:
                     self._replica_id.unique_id, self._node_id
                 )
                 self._ready_obj_ref = replica_ready_check_func.remote(
-                    deployment_config, self._rank
+                    deployment_config, self._rank, self._gang_context
                 )
 
             return ReplicaStartupStatus.PENDING_INITIALIZATION, None
@@ -1376,6 +1384,7 @@ def start(
         assign_rank_callback: Callable[[ReplicaID], ReplicaRank],
         gang_placement_group: Optional[PlacementGroup] = None,
         gang_bundle_index: Optional[int] = None,
+        gang_context: Optional[GangContext] = None,
     ) -> ReplicaSchedulingRequest:
         """
         Start a new actor for current DeploymentReplica instance.
@@ -1385,12 +1394,14 @@ def start(
             assign_rank_callback: Callback to assign rank to the replica.
             gang_placement_group: Pre-created gang PG to schedule this replica on.
             gang_bundle_index: Bundle index within the gang PG for this replica.
+            gang_context: Gang context for this replica (if part of a gang).
         """
         replica_scheduling_request = self._actor.start(
             deployment_info,
             assign_rank_callback=assign_rank_callback,
             gang_placement_group=gang_placement_group,
             gang_bundle_index=gang_bundle_index,
+            gang_context=gang_context,
         )
         self._start_time = time.time()
         self._logged_shutdown_message = False
@@ -1431,6 +1442,11 @@ def rank(self) -> Optional[ReplicaRank]:
         """Get the rank assigned to the replica."""
         return self._actor.rank
 
+    @property
+    def gang_context(self) -> Optional[GangContext]:
+        """Get the gang context for this replica (if part of a gang)."""
+        return self._actor.gang_context
+
     def check_started(
         self,
     ) -> Tuple[ReplicaStartupStatus, Optional[str], Optional[float]]:
@@ -3060,20 +3076,35 @@ def _add_replicas_with_gang_scheduling(
                     "This should not happen."
                 )
 
-            for bundle_index in range(gang_size):
-                replica_id = ReplicaID(get_random_string(), deployment_id=self._id)
+            # Pre-generate replica IDs for all members of this gang
+            gang_id = get_random_string()
+            member_replica_ids = [
+                ReplicaID(get_random_string(), deployment_id=self._id)
+                for _ in range(gang_size)
+            ]
+
+            for bundle_index, replica_id in enumerate(member_replica_ids):
+                gang_context = GangContext(
+                    gang_id=gang_id,
+                    rank=bundle_index,
+                    world_size=gang_size,
+                    member_replica_ids=[
+                        r.unique_id for r in member_replica_ids
+                    ],
+                )
 
                 new_deployment_replica = DeploymentReplica(
                     replica_id,
                     self._target_state.version,
                 )
 
-                # Start the replica with gang PG information
+                # Start the replica with gang PG and gang context
                 scheduling_request = new_deployment_replica.start(
                     self._target_state.info,
                     assign_rank_callback=self._rank_manager.assign_rank,
                     gang_placement_group=gang_pg,
                     gang_bundle_index=bundle_index,
+                    gang_context=gang_context,
                 )
 
                 upscale.append(scheduling_request)

python/ray/serve/_private/replica.py

@@ -934,6 +934,9 @@ def __init__(
         # Flipped to `True` once graceful shutdown is initiated. May be used by replica
         # subclass implementations.
         self._shutting_down = False
+        # Gang context for this replica (if part of a gang).
+        # Populated during initialize_and_get_metadata.
+        self._gang_context: Optional["GangContext"] = None
 
         # Will be populated with the wrapped ASGI app if the user callable is an
         # `ASGIAppReplicaWrapper` (i.e., they are using the FastAPI integration).
@@ -1070,6 +1073,7 @@ def register_handle_callback(deployment_id: DeploymentID) -> None:
             rank=rank,
             world_size=world_size,
             handle_registration_callback=register_handle_callback,
+            gang_context=self._gang_context,
         )
 
     def _configure_logger_and_profilers(
@@ -1382,8 +1386,13 @@ async def _on_initialized(self):
         raise NotImplementedError
 
     async def initialize(
-        self, deployment_config: Optional[DeploymentConfig], rank: Optional[ReplicaRank]
+        self,
+        deployment_config: Optional[DeploymentConfig],
+        rank: Optional[ReplicaRank],
+        gang_context: Optional["GangContext"] = None,
     ):
+        if gang_context is not None:
+            self._gang_context = gang_context
         if rank is not None:
             self._rank = rank
             self._set_internal_replica_context(
@@ -2427,7 +2436,10 @@ def list_outbound_deployments(self) -> Optional[List[DeploymentID]]:
         return self._replica_impl.list_outbound_deployments()
 
     async def initialize_and_get_metadata(
-        self, deployment_config: DeploymentConfig = None, rank: ReplicaRank = None
+        self,
+        deployment_config: DeploymentConfig = None,
+        rank: ReplicaRank = None,
+        gang_context: "GangContext" = None,
     ) -> ReplicaMetadata:
         """Handles initializing the replica.
 
@@ -2440,7 +2452,7 @@ async def initialize_and_get_metadata(
         """
         # Unused `_after` argument is for scheduling: passing an ObjectRef
         # allows delaying this call until after the `_after` call has returned.
-        await self._replica_impl.initialize(deployment_config, rank)
+        await self._replica_impl.initialize(deployment_config, rank, gang_context)
         return self._replica_impl.get_metadata()
 
     async def check_health(self):

python/ray/serve/context.py

@@ -13,7 +13,7 @@
 import ray
 from ray.exceptions import RayActorError
 from ray.serve._private.client import ServeControllerClient
-from ray.serve._private.common import DeploymentID, ReplicaID
+from ray.serve._private.common import DeploymentID, GangContext, ReplicaID
 from ray.serve._private.config import DeploymentConfig
 from ray.serve._private.constants import (
     SERVE_CONTROLLER_NAME,
@@ -24,32 +24,14 @@
 from ray.serve.exceptions import RayServeException
 from ray.serve.grpc_util import RayServegRPCContext
 from ray.serve.schema import ReplicaRank
-from ray.util.annotations import DeveloperAPI, PublicAPI
+from ray.util.annotations import DeveloperAPI
 
 logger = logging.getLogger(SERVE_LOGGER_NAME)
 
 _INTERNAL_REPLICA_CONTEXT: "ReplicaContext" = None
 _global_client: ServeControllerClient = None
 
 
-@PublicAPI(stability="alpha")
-@dataclass
-class GangContext:
-    """Context information for a replica that is part of a gang.
-
-    Attributes:
-        gang_id: Unique identifier for this gang.
-        rank: This replica's rank within the gang (0-indexed).
-        world_size: Total number of replicas in this gang.
-        member_replica_ids: List of replica IDs in this gang, ordered by rank.
-    """
-
-    gang_id: str
-    rank: int
-    world_size: int
-    member_replica_ids: List[str]
-
-
 @DeveloperAPI
 @dataclass
 class ReplicaContext:

python/ray/serve/tests/test_deployment_scheduler.py

@@ -736,8 +736,7 @@ class TestGangScheduling:
     """Tests for gang scheduling with placement groups."""
 
     def test_sufficient_resources(self, ray_start_cluster):
-        """Verifies that gang scheduling succeeds when cluster has sufficient resources.
-        """
+        """Verifies that gang scheduling succeeds when cluster has sufficient resources."""
         cluster = ray_start_cluster
         cluster.add_node(num_cpus=4)
         cluster.add_node(num_cpus=4)
@@ -771,8 +770,7 @@ def GangDeployment():
         serve.shutdown()
 
     def test_insufficient_resources(self, ray_start_cluster):
-        """Verifies that gang scheduling fails when cluster lacks resources.
-        """
+        """Verifies that gang scheduling fails when cluster lacks resources."""
         cluster = ray_start_cluster
         cluster.add_node(num_cpus=4)
         cluster.add_node(num_cpus=4)
@@ -807,8 +805,7 @@ def check_deploy_failed():
         serve.shutdown()
 
     def test_pack_strategy(self, ray_start_cluster):
-        """Verifies that PACK strategy places gang replicas on the same node.
-        """
+        """Verifies that PACK strategy places gang replicas on the same node."""
         cluster = ray_start_cluster
         cluster.add_node(num_cpus=4)
         cluster.add_node(num_cpus=4)
@@ -895,5 +892,97 @@ def SpreadDeployment():
         serve.shutdown()
 
 
+    def test_gang_context_populated(self, ray_start_cluster):
+        """Verifies GangContext is correctly populated in ReplicaContext."""
+        cluster = ray_start_cluster
+        cluster.add_node(num_cpus=4)
+        cluster.wait_for_nodes()
+        ray.init(address=cluster.address)
+        serve.start()
+
+        @serve.deployment
+        class GangContextDeployment:
+            def __call__(self):
+                ctx = ray.serve.context._get_internal_replica_context()
+                gc = ctx.gang_context
+                if gc is None:
+                    return None
+                return {
+                    "gang_id": gc.gang_id,
+                    "rank": gc.rank,
+                    "world_size": gc.world_size,
+                    "member_replica_ids": gc.member_replica_ids,
+                    "replica_id": ctx.replica_id.unique_id,
+                }
+
+        app = GangContextDeployment.options(
+            num_replicas=4,
+            ray_actor_options={"num_cpus": 1},
+            gang_scheduling_config=GangSchedulingConfig(gang_size=2),
+        ).bind()
+
+        handle = serve.run(app, name="gang_context_app")
+        wait_for_condition(
+            check_apps_running,
+            apps=["gang_context_app"],
+            timeout=60,
+        )
+
+        # Collect gang contexts from all replicas
+        # Query enough times to hit all 4 replicas
+        contexts_by_replica = {}
+        for _ in range(100):
+            result = handle.remote().result()
+            assert result is not None
+            replica_id = result["replica_id"]
+            if replica_id not in contexts_by_replica:
+                contexts_by_replica[replica_id] = result
+            if len(contexts_by_replica) == 4:
+                break
+        assert len(contexts_by_replica) == 4
+
+        # Group replicas by gang_id
+        gangs = {}
+        for replica_id, ctx in contexts_by_replica.items():
+            gang_id = ctx["gang_id"]
+            gangs.setdefault(gang_id, []).append(ctx)
+
+        # Should have exactly 2 gangs
+        assert len(gangs) == 2
+
+        for gang_id, members in gangs.items():
+            # Each gang should have exactly 2 replicas
+            assert len(members) == 2
+
+            # All members should have the same world_size
+            assert all(member["world_size"] == 2 for member in members)
+
+            # All members should have the same member_replica_ids
+            assert members[0]["member_replica_ids"] == members[1]["member_replica_ids"]
+
+            # member_replica_ids should contain exactly the 2 replica IDs in this gang
+            expected_ids = sorted([m["replica_id"] for m in members])
+            actual_ids = sorted(members[0]["member_replica_ids"])
+            assert actual_ids == expected_ids
+
+            # Ranks within the gang should be {0, 1}
+            ranks = sorted([m["rank"] for m in members])
+            assert ranks == [0, 1]
+
+        # Across gangs: gang_ids should be different (already guaranteed by dict keys)
+        gang_ids = list(gangs.keys())
+        assert gang_ids[0] != gang_ids[1]
+
+        # Across gangs: member_replica_ids should be different
+        gang_members_list = list(gangs.values())
+        assert (
+            sorted(gang_members_list[0][0]["member_replica_ids"])
+            != sorted(gang_members_list[1][0]["member_replica_ids"])
+        )
+
+        serve.delete("gang_context_app")
+        serve.shutdown()
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main(["-v", "-s", __file__]))