[serve] Cache per-tick replica id walks behind a container mutation version

A CPU profile of a controller managing 16K replicas showed the loop
dominated by O(N) id-collection walks recomputed every tick (alive
actor ids, active node ids, running replica ids, the autoscaler id-set
rebuild) even though membership rarely changes in steady state.

ReplicaStateContainer now bumps a mutation version on add / non-empty
pop / non-empty remove, and the derived collections memoize on it at
both DeploymentState and DeploymentStateManager level. Caching disarms
while STARTING/RECOVERING replicas exist since their actor/node ids
materialize in place. The autoscaler skips its id-set rebuild when
handed the same (cached) list object, and the rank-consistency
fingerprint becomes the version integer.

Benchmark at 16K replicas: control loop 352.6ms -> 42.9ms.

Signed-off-by: john.taylor <john.taylor@anyscale.com>

Author: johntaylor-cell

python/ray/serve/_private/autoscaling_state.py

@@ -210,6 +210,9 @@ def get_num_replicas_upper_bound(self) -> int:
 
     def update_running_replica_ids(self, running_replicas: List[ReplicaID]):
         """Update cached set of running replica IDs for this deployment."""
+        if running_replicas is self._running_replicas:
+            # Same (cached) list object -- membership unchanged, skip rebuild.
+            return
         self._running_replicas = running_replicas
         self._cached_running_replica_strs = {
             r.to_full_id_str() for r in running_replicas

python/ray/serve/_private/deployment_state.py

@@ -10,8 +10,6 @@
 from copy import copy
 from dataclasses import dataclass
 from enum import Enum
-from functools import reduce
-from operator import xor
 from typing import Any, Callable, Deque, Dict, List, Optional, Set, Tuple
 
 import ray
@@ -2159,6 +2157,20 @@ def get_outbound_deployments(self) -> Optional[List[DeploymentID]]:
         return self._actor.get_outbound_deployments()
 
 
+def _memoized_walk(obj, memo_attr: str, token, compute):
+    """Return compute() memoized on *token* in obj.<memo_attr>.
+
+    token=None means "uncacheable right now": recompute and drop the memo.
+    Callers must treat returned collections as immutable (they are shared).
+    """
+    memo = getattr(obj, memo_attr)
+    if token is not None and memo is not None and memo[0] == token:
+        return memo[1]
+    result = compute()
+    setattr(obj, memo_attr, None if token is None else (token, result))
+    return result
+
+
 class ReplicaStateContainer:
     """Container for mapping ReplicaStates to lists of DeploymentReplicas."""
 
@@ -2169,6 +2181,12 @@ def __init__(self, on_replica_state_change=None):
         # Incremental (state, version) counts for O(1) version-filtered count()
         # (maintained on add/pop/remove) -> replaces the per-tick O(N) version scan.
         self._sv_counts: Dict[tuple, int] = defaultdict(int)
+        # Bumped on every mutation; derived id collections memoize on it.
+        self._mutation_version: int = 0
+
+    @property
+    def mutation_version(self) -> int:
+        return self._mutation_version
 
     def __getstate__(self):
         # Exclude the callback to keep the container picklable (the callback
@@ -2192,6 +2210,7 @@ def add(self, state: ReplicaState, replica: DeploymentReplica):
         self._replicas[state].append(replica)
         self._replica_id_index[replica.replica_id] = replica
         self._sv_counts[(state, replica.version)] += 1
+        self._mutation_version += 1
         if self._on_replica_state_change and state != old_state:
             self._on_replica_state_change(old_state, state)
 
@@ -2289,6 +2308,8 @@ def pop(
         for replica in replicas:
             self._replica_id_index.pop(replica.replica_id, None)
 
+        if replicas:
+            self._mutation_version += 1
         return replicas
 
     def count(
@@ -2362,6 +2383,8 @@ def remove(self, replica_ids: Set[ReplicaID]) -> List[DeploymentReplica]:
                     remaining.append(replica)
             if found_any:
                 self._replicas[state] = remaining
+        if removed:
+            self._mutation_version += 1
         return removed
 
     def __str__(self):
@@ -2922,6 +2945,10 @@ def __init__(
             fail_on_rank_error=RAY_SERVE_FAIL_ON_RANK_ERROR
         )
         self._last_rank_membership_fingerprint: Optional[int] = None
+        # (token, result) memos for per-tick id-collection walks.
+        self._alive_replica_actor_ids_memo = None
+        self._running_replica_ids_memo = None
+        self._active_node_ids_memo = None
 
         self.replica_average_ongoing_requests: Dict[str, float] = {}
 
@@ -3298,16 +3325,42 @@ def _terminally_failed(self) -> bool:
         )
         return replica_failed or self.deployment_actor_terminally_failed()
 
+    def _membership_cache_token(self) -> Optional[int]:
+        """Cache key for derived id collections, or None while uncacheable.
+
+        The container version tracks add/pop/remove, but STARTING/RECOVERING
+        replicas can have actor_id/actor_node_id materialize in place without
+        a container mutation -- disable caching while any exist.
+        """
+        if (
+            self._replicas.count(
+                states=[ReplicaState.STARTING, ReplicaState.RECOVERING]
+            )
+            > 0
+        ):
+            return None
+        return self._replicas.mutation_version
+
     def get_alive_replica_actor_ids(self) -> Set[str]:
-        return {replica.actor_id for replica in self._replicas.get()}
+        return _memoized_walk(
+            self,
+            "_alive_replica_actor_ids_memo",
+            self._membership_cache_token(),
+            lambda: {replica.actor_id for replica in self._replicas.get()},
+        )
 
     def get_running_replica_ids(self) -> List[ReplicaID]:
-        return [
-            replica.replica_id
-            for replica in self._replicas.get(
-                [ReplicaState.RUNNING, ReplicaState.PENDING_MIGRATION]
-            )
-        ]
+        return _memoized_walk(
+            self,
+            "_running_replica_ids_memo",
+            self._membership_cache_token(),
+            lambda: [
+                replica.replica_id
+                for replica in self._replicas.get(
+                    [ReplicaState.RUNNING, ReplicaState.PENDING_MIGRATION]
+                )
+            ],
+        )
 
     def get_running_replica_infos(self) -> List[RunningReplicaInfo]:
         return [
@@ -3359,11 +3412,16 @@ def get_active_node_ids(self) -> Set[str]:
             # node before all the replicas are migrated.
             ReplicaState.PENDING_MIGRATION,
         ]
-        return {
-            replica.actor_node_id
-            for replica in self._replicas.get(active_states)
-            if replica.actor_node_id is not None
-        }
+        return _memoized_walk(
+            self,
+            "_active_node_ids_memo",
+            self._membership_cache_token(),
+            lambda: {
+                replica.actor_node_id
+                for replica in self._replicas.get(active_states)
+                if replica.actor_node_id is not None
+            },
+        )
 
     def list_replica_details(self) -> List[ReplicaDetails]:
         return [replica.actor_details for replica in self._replicas.get()]
@@ -4871,24 +4929,21 @@ def _maybe_check_rank_consistency(self) -> None:
         the controller loop. Rank consistency can only be violated by
         membership changes, so a cheap fingerprint gates it.
         """
-        active_replicas = self._replicas.get()
         if not (
-            active_replicas
-            and self._curr_status_info.status == DeploymentStatus.HEALTHY
+            self._curr_status_info.status == DeploymentStatus.HEALTHY
             # Skip consistency check if there are STARTING replicas. During node
             # migration, new replicas are created in STARTING state (without ranks)
             # after the status is set to HEALTHY. Running the consistency check
             # with STARTING replicas causes "active keys without ranks" error.
             and self._replicas.count(states=[ReplicaState.STARTING]) == 0
         ):
             return
-        fingerprint = reduce(
-            xor,
-            (hash(r.replica_id.unique_id) for r in active_replicas),
-            len(active_replicas),
-        )
+        fingerprint = self._replicas.mutation_version
         if fingerprint == self._last_rank_membership_fingerprint:
             return
+        active_replicas = self._replicas.get()
+        if not active_replicas:
+            return
         replicas_to_reconfigure = (
             self._rank_manager.check_rank_consistency_and_reassign_minimally(
                 active_replicas,
@@ -5665,6 +5720,9 @@ def __init__(
         self._shutting_down = False
 
         self._deployment_states: Dict[DeploymentID, DeploymentState] = {}
+        # (key, result) memos for cross-deployment id-collection walks.
+        self._alive_replica_actor_ids_memo = None
+        self._active_node_ids_memo = None
         # Monotonic counter bumped whenever an ingress deployment's running-replica
         # set (node/ports included) changes; the controller gates the direct-ingress port
         # reconcile on it, skipping the O(replicas) pass on ticks with no change.
@@ -5991,12 +6049,30 @@ def get_deployment_statuses(
                     statuses.append(state.curr_status_info)
             return statuses
 
+    def _membership_cache_key(self) -> Optional[tuple]:
+        """Combined per-deployment token, or None if any deployment is
+        uncacheable. Deployment add/remove changes the key shape."""
+        parts = []
+        for deployment_id, ds in self._deployment_states.items():
+            token = ds._membership_cache_token()
+            if token is None:
+                return None
+            parts.append((deployment_id, token))
+        return tuple(parts)
+
     def get_alive_replica_actor_ids(self) -> Set[str]:
-        alive_replica_actor_ids = set()
-        for ds in self._deployment_states.values():
-            alive_replica_actor_ids |= ds.get_alive_replica_actor_ids()
+        def compute():
+            alive_replica_actor_ids = set()
+            for ds in self._deployment_states.values():
+                alive_replica_actor_ids |= ds.get_alive_replica_actor_ids()
+            return alive_replica_actor_ids
 
-        return alive_replica_actor_ids
+        return _memoized_walk(
+            self,
+            "_alive_replica_actor_ids_memo",
+            self._membership_cache_key(),
+            compute,
+        )
 
     def get_deployment_ids(self) -> List[DeploymentID]:
         return list(self._deployment_states.keys())
@@ -6424,10 +6500,16 @@ def get_active_node_ids(self) -> Set[str]:
         This is used to determine which node has replicas. Only nodes with replicas and
         head node should have active proxies.
         """
-        node_ids = set()
-        for deployment_state in self._deployment_states.values():
-            node_ids.update(deployment_state.get_active_node_ids())
-        return node_ids
+
+        def compute():
+            node_ids = set()
+            for deployment_state in self._deployment_states.values():
+                node_ids.update(deployment_state.get_active_node_ids())
+            return node_ids
+
+        return _memoized_walk(
+            self, "_active_node_ids_memo", self._membership_cache_key(), compute
+        )
 
     def get_ingress_membership_version(self) -> int:
         """Monotonic counter of ingress running-replica-set changes (node/ports

python/ray/serve/tests/unit/test_deployment_state.py

@@ -9975,6 +9975,7 @@ def _ds(self, uids):
         ds._replicas = Mock()
         ds._replicas.get.return_value = replicas
         ds._replicas.count.return_value = 0
+        ds._replicas.mutation_version = 7
         ds._curr_status_info = Mock(status=DeploymentStatus.HEALTHY)
         ds._rank_manager = Mock()
         ds._rank_manager.check_rank_consistency_and_reassign_minimally.return_value = []
@@ -10001,6 +10002,8 @@ def test_membership_change_reruns(self):
             r.replica_id.unique_id = uid
             new.append(r)
         ds._replicas.get.return_value = new
+        # Membership changes always bump the container mutation version.
+        ds._replicas.mutation_version = 8
         ds._maybe_check_rank_consistency()
         assert (
             ds._rank_manager.check_rank_consistency_and_reassign_minimally.call_count
@@ -10015,5 +10018,81 @@ def test_starting_replicas_skip_entirely(self):
         assert ds._last_rank_membership_fingerprint is None
 
 
+class TestMutationVersionMemo:
+    """Container mutation version + memoized per-tick id-collection walks."""
+
+    def _replica(self, uid, actor_id=None, node_id=None):
+        r = Mock()
+        r.replica_id = uid
+        r.version = "v1"
+        r.actor_id = actor_id or f"actor-{uid}"
+        r.actor_node_id = node_id or f"node-{uid}"
+        return r
+
+    def test_version_bumps_only_on_real_mutation(self):
+        c = ds_mod.ReplicaStateContainer()
+        v0 = c.mutation_version
+        r = self._replica("a")
+        c.add(ReplicaState.STARTING, r)
+        assert c.mutation_version == v0 + 1
+        # Empty pops/removes (the steady-state case) do not bump.
+        assert c.pop(states=[ReplicaState.STOPPING]) == []
+        assert c.remove({"missing"}) == []
+        assert c.mutation_version == v0 + 1
+        assert c.pop(states=[ReplicaState.STARTING]) == [r]
+        assert c.mutation_version == v0 + 2
+
+    def _ds(self):
+        ds = ds_mod.DeploymentState.__new__(ds_mod.DeploymentState)
+        ds._replicas = ds_mod.ReplicaStateContainer()
+        ds._alive_replica_actor_ids_memo = None
+        ds._running_replica_ids_memo = None
+        ds._active_node_ids_memo = None
+        return ds
+
+    def test_getters_memoize_and_invalidate_on_mutation(self):
+        ds = self._ds()
+        ds._replicas.add(ReplicaState.RUNNING, self._replica("r1"))
+        ids1 = ds.get_running_replica_ids()
+        assert ids1 == ["r1"]
+        assert ds.get_running_replica_ids() is ids1  # cache hit
+        assert ds.get_alive_replica_actor_ids() == {"actor-r1"}
+        assert ds.get_active_node_ids() == {"node-r1"}
+
+        ds._replicas.add(ReplicaState.RUNNING, self._replica("r2"))
+        ids2 = ds.get_running_replica_ids()
+        assert set(ids2) == {"r1", "r2"} and ids2 is not ids1
+        assert ds.get_alive_replica_actor_ids() == {"actor-r1", "actor-r2"}
+        assert ds.get_active_node_ids() == {"node-r1", "node-r2"}
+
+    def test_starting_or_recovering_disables_caching(self):
+        ds = self._ds()
+        ds._replicas.add(ReplicaState.RUNNING, self._replica("r1"))
+        ds._replicas.add(ReplicaState.STARTING, self._replica("r2"))
+        a = ds.get_alive_replica_actor_ids()
+        b = ds.get_alive_replica_actor_ids()
+        # Recomputed each call (attrs may materialize in place), never cached.
+        assert a == b and a is not b
+        assert ds._alive_replica_actor_ids_memo is None
+
+    def test_autoscaling_update_skips_same_list_object(self):
+        from ray.serve._private import autoscaling_state as as_mod
+
+        das = as_mod.DeploymentAutoscalingState.__new__(
+            as_mod.DeploymentAutoscalingState
+        )
+        das._running_replicas = []
+        rid = Mock()
+        rid.to_full_id_str.return_value = "d#r1"
+        ids = [rid]
+        das.update_running_replica_ids(ids)
+        assert das._cached_running_replica_strs == {"d#r1"}
+        das._cached_running_replica_strs = {"sentinel"}
+        das.update_running_replica_ids(ids)  # same object -> skipped
+        assert das._cached_running_replica_strs == {"sentinel"}
+        das.update_running_replica_ids(list(ids))  # new object -> rebuilt
+        assert das._cached_running_replica_strs == {"d#r1"}
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main(["-v", "-s", __file__]))