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

Squashed for a clean rebase onto the A6 tip.

The per-tick id walks (alive-actor-ids, active-node-ids, running-ids, the
autoscaler id-set rebuild) recompute identical results every steady-state tick.
Memoize each on a ReplicaStateContainer mutation version, bumped on add /
non-empty pop / remove, and disarm while STARTING or RECOVERING replicas exist.

Review fixes included: manager-level getters return frozensets so the memo
cannot be mutated by callers (with set() at the one mutating caller),
AutoscalingContext.running_replicas is copied to a list at the public boundary,
annotations corrected, and an end-to-end test that 20 steady-state ticks do not
bump the version and all three getters return identical objects.

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
@@ -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,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, 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
+
+
 class ReplicaStateContainer:
     """Container for mapping ReplicaStates to lists of DeploymentReplicas."""
 
@@ -2181,6 +2205,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 getattr(self, "_mutation_version", 0)
 
     def __getstate__(self):
         # Exclude the callback to keep the container picklable (the callback
@@ -2204,6 +2234,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 = getattr(self, "_mutation_version", 0) + 1
         if self._on_replica_state_change and state != old_state:
             self._on_replica_state_change(old_state, state)
 
@@ -2301,6 +2332,8 @@ def pop(
         for replica in replicas:
             self._replica_id_index.pop(replica.replica_id, None)
 
+        if replicas:
+            self._mutation_version = getattr(self, "_mutation_version", 0) + 1
         return replicas
 
     def count(
@@ -2374,6 +2407,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 = getattr(self, "_mutation_version", 0) + 1
         return removed
 
     def __str__(self):
@@ -2950,7 +2985,11 @@ 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
+        self._last_rank_membership_version: 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] = {}
 
@@ -3327,16 +3366,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 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) -> 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 [
@@ -3388,11 +3453,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()]
@@ -4997,7 +5067,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 mutation version 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.
@@ -5010,11 +5080,12 @@ def _maybe_check_rank_consistency(self) -> None:
             or self._replicas.count(states=[ReplicaState.STARTING]) != 0
         ):
             return
+        version = self._replicas.mutation_version
+        if version == getattr(self, "_last_rank_membership_version", None):
+            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_version = version
             return
         replicas_to_reconfigure = (
             self._rank_manager.check_rank_consistency_and_reassign_minimally(
@@ -5030,7 +5101,7 @@ def _maybe_check_rank_consistency(self) -> None:
         if checked_cleanly:
             # Deliberate: a deployment that keeps erroring never caches and is
             # rechecked every tick, rather than latching an unvalidated membership.
-            self._last_rank_membership_ids = active_replica_ids
+            self._last_rank_membership_version = version
 
     def _handle_deployment_actor_failed_health_check(
         self,
@@ -5798,6 +5869,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.
@@ -6124,12 +6198,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())
@@ -6556,16 +6649,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

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

@@ -10281,5 +10281,34 @@ def test_starting_replicas_skip_the_pass(
         assert ds._rank_manager.consistency_calls == 0
 
 
+def test_steady_state_ticks_do_not_bump_mutation_version(mock_deployment_state_manager):
+    """The memo's whole value: an idle tick must not mutate the container."""
+    create_dsm, _, _, _ = mock_deployment_state_manager
+    dsm = create_dsm()
+    info, v = deployment_info(num_replicas=8, version="1")
+    assert dsm.deploy(TEST_DEPLOYMENT_ID, info)
+    dsm.save_checkpoint()
+    ds = dsm._get_deployment_state_for_testing(TEST_DEPLOYMENT_ID)
+    dsm.update()
+    for replica in ds._replicas.get([ReplicaState.STARTING]):
+        replica._actor.set_ready()
+    dsm.update()
+    check_counts(ds, total=8, by_state=[(ReplicaState.RUNNING, 8, v)])
+    for _ in range(3):
+        dsm.update()
+    v0 = ds._replicas.mutation_version
+    alive, running, nodes = (
+        ds.get_alive_replica_actor_ids(),
+        ds.get_running_replica_ids(),
+        ds.get_active_node_ids(),
+    )
+    for _ in range(20):
+        dsm.update()
+        assert ds._replicas.mutation_version == v0
+        assert ds.get_alive_replica_actor_ids() is alive
+        assert ds.get_running_replica_ids() is running
+        assert ds.get_active_node_ids() is nodes
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main(["-v", "-s", __file__]))