[serve] Cache per-tick replica id walks behind a container membership fingerprint

The per-tick id walks (alive-actor-ids, active-node-ids, running-ids, the
autoscaler id-set rebuild) recompute identical results on every steady-state
tick. Memoize each on a ReplicaStateContainer membership fingerprint, and disarm
while STARTING or RECOVERING replicas exist, since their actor_id/actor_node_id
materialize in place without touching the container.

The fingerprint is the summed per-replica hash of the container's
{replica id -> state} content rather than a mutation counter, because the health
and migration passes pop a whole bucket and re-add every replica to the state it
came from. A counter reads that churn as a membership change, and one node
draining anywhere in the cluster is enough to trigger it: at 4 replicas it bumps
5 times per otherwise-idle tick, which invalidates all three memos and re-runs
the O(N) rank-consistency pass on every tick. Summing per-replica terms makes
pop-and-re-add cancel exactly, while a real arrival, departure or state
transition still changes it.

The same fingerprint replaces the id-set key of the rank-consistency gate and of
its error backoff, so both keep an O(1) key without losing the content-based
semantics that made them hold under churn.

Manager-level getters return frozensets so the memo cannot be mutated by callers
(with set() at the one mutating caller), and AutoscalingContext.running_replicas
is copied to a list at the public boundary.

Signed-off-by: john.taylor <john.taylor@anyscale.com>
Co-Authored-By: Claude <noreply@anthropic.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
@@ -395,7 +398,9 @@ def get_autoscaling_context(
             app_name=self._deployment_id.app_name,
             current_num_replicas=len(self._running_replicas),
             target_num_replicas=curr_target_num_replicas,
-            running_replicas=self._running_replicas,
+            # list(): AutoscalingContext.running_replicas is List[ReplicaID] on a
+            # stable PublicAPI; the tuple stays internal so the identity skip works.
+            running_replicas=list(self._running_replicas),
             total_num_requests=self.get_total_num_requests,
             capacity_adjusted_min_replicas=self.get_num_replicas_lower_bound(),
             capacity_adjusted_max_replicas=self.get_num_replicas_upper_bound(),

python/ray/serve/_private/controller.py

@@ -533,7 +533,8 @@ def _update_proxy_nodes(self):
         Controller decides where proxy actors should run
         (head node and nodes with deployment replicas).
         """
-        new_proxy_nodes = self.deployment_state_manager.get_active_node_ids()
+        # set(): the getter returns a frozenset memo, and this is mutated below.
+        new_proxy_nodes = set(self.deployment_state_manager.get_active_node_ids())
         new_proxy_nodes = new_proxy_nodes - set(
             self.cluster_node_info_cache.get_draining_nodes()
         )

python/ray/serve/_private/deployment_state.py

@@ -10,7 +10,17 @@
 from copy import copy
 from dataclasses import dataclass
 from enum import Enum
-from typing import Any, Callable, Deque, Dict, List, Optional, Set, Tuple
+from typing import (
+    Any,
+    Callable,
+    Deque,
+    Dict,
+    FrozenSet,
+    List,
+    Optional,
+    Set,
+    Tuple,
+)
 
 import ray
 from ray import ObjectRef, cloudpickle
@@ -2171,6 +2181,25 @@ 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, None)
+    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
+
+
+def _membership_term(replica_id: ReplicaID, state: ReplicaState) -> int:
+    """One replica's term in ReplicaStateContainer.membership_fingerprint."""
+    return hash((replica_id, state))
+
+
 class ReplicaStateContainer:
     """Container for mapping ReplicaStates to lists of DeploymentReplicas."""
 
@@ -2181,6 +2210,18 @@ 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)
+        # Derived id collections memoize on this (see membership_fingerprint).
+        self._membership_fingerprint: int = 0
+
+    @property
+    def membership_fingerprint(self) -> int:
+        """Fingerprint of the {replica id -> state} content, for memo keys.
+
+        Summed per-replica terms, so the pop-and-re-add churn of the health and
+        migration passes cancels exactly while real arrivals, departures and
+        state transitions change it.
+        """
+        return self._membership_fingerprint
 
     def __getstate__(self):
         # Exclude the callback to keep the container picklable (the callback
@@ -2204,6 +2245,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._membership_fingerprint += _membership_term(replica.replica_id, state)
         if self._on_replica_state_change and state != old_state:
             self._on_replica_state_change(old_state, state)
 
@@ -2297,6 +2339,7 @@ def pop(
             replicas.extend(popped)
             for _r in popped:
                 self._sv_counts[(state, _r.version)] -= 1
+                self._membership_fingerprint -= _membership_term(_r.replica_id, state)
 
         for replica in replicas:
             self._replica_id_index.pop(replica.replica_id, None)
@@ -2367,6 +2410,9 @@ def remove(self, replica_ids: Set[ReplicaID]) -> List[DeploymentReplica]:
                 if remaining_to_find > 0 and replica.replica_id in replica_ids:
                     removed.append(replica)
                     self._sv_counts[(state, replica.version)] -= 1
+                    self._membership_fingerprint -= _membership_term(
+                        replica.replica_id, state
+                    )
                     self._replica_id_index.pop(replica.replica_id, None)
                     remaining_to_find -= 1
                     found_any = True
@@ -2955,10 +3001,14 @@ def __init__(
         self._rank_manager = DeploymentRankManager(
             fail_on_rank_error=RAY_SERVE_FAIL_ON_RANK_ERROR
         )
-        self._last_rank_membership_ids: Optional[Set[str]] = None
-        # Membership whose pass swallowed an error, and when, for retry backoff.
-        self._last_rank_error_ids: Optional[Set[str]] = None
+        self._last_rank_membership_fp: Optional[int] = None
+        # Fingerprint whose pass swallowed an error, and when, for retry backoff.
+        self._last_rank_error_fp: Optional[int] = None
         self._last_rank_error_ts: float = 0.0
+        # (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] = {}
 
@@ -3335,16 +3385,42 @@ def _terminally_failed(self) -> bool:
         )
         return replica_failed or self.deployment_actor_terminally_failed()
 
-    def get_alive_replica_actor_ids(self) -> Set[str]:
-        return {replica.actor_id for replica in self._replicas.get()}
+    def _membership_cache_token(self) -> Optional[int]:
+        """Cache key for derived id collections, or None while uncacheable.
 
-    def get_running_replica_ids(self) -> List[ReplicaID]:
-        return [
-            replica.replica_id
-            for replica in self._replicas.get(
-                [ReplicaState.RUNNING, ReplicaState.PENDING_MIGRATION]
+        The container fingerprint tracks membership, but STARTING/RECOVERING
+        replicas can have actor_id/actor_node_id materialize in place without
+        a membership change -- disable caching while any exist.
+        """
+        if (
+            self._replicas.count(
+                states=[ReplicaState.STARTING, ReplicaState.RECOVERING]
             )
-        ]
+            > 0
+        ):
+            return None
+        return self._replicas.membership_fingerprint
+
+    def get_alive_replica_actor_ids(self) -> FrozenSet[str]:
+        return _memoized_walk(
+            self,
+            "_alive_replica_actor_ids_memo",
+            self._membership_cache_token(),
+            lambda: frozenset(replica.actor_id for replica in self._replicas.get()),
+        )
+
+    def get_running_replica_ids(self) -> Tuple[ReplicaID, ...]:
+        return _memoized_walk(
+            self,
+            "_running_replica_ids_memo",
+            self._membership_cache_token(),
+            lambda: tuple(
+                replica.replica_id
+                for replica in self._replicas.get(
+                    [ReplicaState.RUNNING, ReplicaState.PENDING_MIGRATION]
+                )
+            ),
+        )
 
     def get_running_replica_infos(self) -> List[RunningReplicaInfo]:
         return [
@@ -3396,11 +3472,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: frozenset(
+                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()]
@@ -5005,7 +5086,7 @@ def _maybe_check_rank_consistency(self) -> None:
         The pass is O(N) with heavy constants and used to run on every
         control-loop tick in steady state -- at 10K+ replicas it monopolizes
         the controller loop. Rank consistency can only be violated by
-        membership changes, so the active replica-id set gates it.
+        membership changes, so the container membership fingerprint gates it.
         """
         # O(1) guards first -- `get()` below copies the whole replica list, and these
         # two rule out the busiest ticks (rollouts, migrations) without paying for it.
@@ -5018,17 +5099,18 @@ def _maybe_check_rank_consistency(self) -> None:
             or self._replicas.count(states=[ReplicaState.STARTING]) != 0
         ):
             return
+        fp = self._replicas.membership_fingerprint
+        if fp == self._last_rank_membership_fp:
+            return
         active_replicas = self._replicas.get()
         if not active_replicas:
-            return
-        active_replica_ids = {r.replica_id.unique_id for r in active_replicas}
-        if active_replica_ids == self._last_rank_membership_ids:
+            self._last_rank_membership_fp = fp
             return
         # An errored pass leaves the membership uncached so it is retried; rate-limit that
         # retry while the membership still matches the one that failed. A real membership
         # change bypasses this and runs immediately.
         if (
-            self._last_rank_error_ids == active_replica_ids
+            self._last_rank_error_fp == fp
             and time.time() - self._last_rank_error_ts < _RANK_ERROR_RETRY_S
         ):
             return
@@ -5044,11 +5126,11 @@ def _maybe_check_rank_consistency(self) -> None:
         # Reconfigure replicas that had their ranks reassigned
         self._reconfigure_replicas_with_new_ranks(replicas_to_reconfigure)
         if checked_cleanly:
-            self._last_rank_membership_ids = active_replica_ids
-            self._last_rank_error_ids = None
+            self._last_rank_membership_fp = fp
+            self._last_rank_error_fp = None
         else:
             # Not cached: an unvalidated membership must be rechecked, just not every tick.
-            self._last_rank_error_ids = active_replica_ids
+            self._last_rank_error_fp = fp
             self._last_rank_error_ts = time.time()
 
     def _handle_deployment_actor_failed_health_check(
@@ -5817,6 +5899,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.
@@ -6143,12 +6228,31 @@ def get_deployment_statuses(
                     statuses.append(state.curr_status_info)
             return statuses
 
-    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()
-
-        return alive_replica_actor_ids
+    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) -> FrozenSet[str]:
+        def compute():
+            alive_replica_actor_ids = set()
+            for ds in self._deployment_states.values():
+                alive_replica_actor_ids |= ds.get_alive_replica_actor_ids()
+            # Frozen: the memo is handed to callers and must not be mutable.
+            return frozenset(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())
@@ -6575,16 +6679,23 @@ def record_request_routing_info(self, info: RequestRoutingInfo) -> None:
             return
         self._deployment_states[deployment_id].record_request_routing_info(info)
 
-    def get_active_node_ids(self) -> Set[str]:
+    def get_active_node_ids(self) -> FrozenSet[str]:
         """Return set of node ids with running replicas of any deployment.
 
         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())
+            # Frozen: the memo is handed to callers and must not be mutable.
+            return frozenset(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