[rollout] feat: re-evaluate sticky sessions against replica load

tests/workers/rollout/test_sticky_affinity_on_cpu.py

@@ -0,0 +1,164 @@
+# Copyright 2026 Amazon.com Inc and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Bounded sticky affinity in ``GlobalRequestLoadBalancer``.
+
+The balancer is a plain class wrapped with ``ray.remote`` at instantiation, so these tests
+drive it directly -- no actor, no GPU.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+
+import pytest
+
+from verl.workers.rollout.llm_server import GlobalRequestLoadBalancer
+
+SERVERS = {"s0": None, "s1": None, "s2": None}
+
+
+def lb(**kwargs) -> GlobalRequestLoadBalancer:
+    kwargs.setdefault("servers", dict(SERVERS))
+    return GlobalRequestLoadBalancer(**kwargs)
+
+
+class TestDefaultBehaviorUnchanged:
+    """The default must be indistinguishable from the pre-change balancer."""
+
+    def test_least_loaded_round_robins(self):
+        balancer = lb()
+        picks = [balancer.acquire_server(f"r{i}")[0] for i in range(6)]
+        assert sorted(picks) == ["s0", "s0", "s1", "s1", "s2", "s2"]
+
+    def test_ties_resolve_to_first_candidate(self):
+        assert lb().acquire_server("r0")[0] == "s0"
+
+    def test_sticky_session_is_honored(self):
+        balancer = lb()
+        pinned = balancer.acquire_server("conv")[0]
+        for _ in range(8):
+            assert balancer.acquire_server("conv")[0] == pinned
+
+    def test_affinity_never_breaks_by_default(self):
+        balancer = lb()
+        pinned = balancer.acquire_server("conv")[0]
+        balancer._inflight_requests[pinned] += 1000
+        assert balancer.acquire_server("conv")[0] == pinned
+        assert balancer.get_status()["affinity_broken"] == 0
+
+    def test_removed_server_still_reselects(self):
+        balancer = lb()
+        pinned = balancer.acquire_server("conv")[0]
+        balancer.remove_servers([pinned])
+        assert balancer.acquire_server("conv")[0] != pinned
+
+    def test_no_servers_raises(self):
+        with pytest.raises(RuntimeError, match="No available servers"):
+            lb(servers={}).acquire_server("r0")
+
+    def test_has_no_async_methods(self):
+        # Any async method makes Ray create an asyncio actor, dropping the FIFO ordering
+        # callers see today.
+        coros = [
+            name
+            for name, fn in inspect.getmembers(GlobalRequestLoadBalancer, inspect.isfunction)
+            if asyncio.iscoroutinefunction(fn)
+        ]
+        assert coros == [], coros
+
+
+class TestBoundedAffinity:
+    def test_overloaded_pin_is_abandoned(self):
+        balancer = lb(affinity_break_margin=2)
+        pinned = balancer.acquire_server("conv")[0]
+        balancer._inflight_requests[pinned] += 10
+        assert balancer.acquire_server("conv")[0] != pinned
+        assert balancer.get_status()["affinity_broken"] == 1
+
+    def test_pin_survives_a_gap_inside_the_margin(self):
+        balancer = lb(affinity_break_margin=5)
+        pinned = balancer.acquire_server("conv")[0]
+        balancer._inflight_requests[pinned] += 4
+        assert balancer.acquire_server("conv")[0] == pinned
+
+    def test_boundary_is_inclusive(self):
+        balancer = lb(affinity_break_margin=3)
+        pinned = balancer.acquire_server("conv")[0]
+        balancer._inflight_requests[pinned] = min(balancer._inflight_requests.values()) + 3
+        assert balancer.acquire_server("conv")[0] == pinned
+        balancer._inflight_requests[pinned] += 1
+        assert balancer.acquire_server("conv")[0] != pinned
+
+    def test_reroute_repins_so_later_turns_follow(self):
+        balancer = lb(affinity_break_margin=0)
+        balancer.acquire_server("conv")
+        balancer._inflight_requests["s0"] = 50
+        balancer._inflight_requests["s1"] = 0
+        balancer._inflight_requests["s2"] = 50
+        assert balancer.acquire_server("conv")[0] == "s1"
+        # The new pin must hold, or every turn re-prefills somewhere new.
+        balancer._inflight_requests["s1"] = 0
+        assert balancer.acquire_server("conv")[0] == "s1"
+
+    def test_zero_margin_tracks_the_least_loaded(self):
+        balancer = lb(affinity_break_margin=0)
+        balancer.acquire_server("conv")
+        for target in ("s2", "s0", "s1"):
+            for sid in balancer._inflight_requests:
+                balancer._inflight_requests[sid] = 0 if sid == target else 9
+            assert balancer.acquire_server("conv")[0] == target
+
+    def test_first_placement_is_not_an_affinity_decision(self):
+        balancer = lb(affinity_break_margin=2)
+        balancer.acquire_server("conv")
+        status = balancer.get_status()
+        assert (status["affinity_kept"], status["affinity_broken"]) == (0, 0)
+        balancer.acquire_server("conv")
+        status = balancer.get_status()
+        assert (status["affinity_kept"], status["affinity_broken"]) == (1, 0)
+
+    def test_full_determinism_ignores_load(self):
+        balancer = lb(affinity_break_margin=0, full_determinism=True)
+        first = balancer.acquire_server("conv")[0]
+        balancer._inflight_requests[first] += 100
+        assert balancer.acquire_server("conv")[0] == first
+        assert balancer.get_status()["affinity_broken"] == 0
+
+    def test_release_lets_a_pin_become_acceptable_again(self):
+        balancer = lb(affinity_break_margin=2)
+        pinned = balancer.acquire_server("conv")[0]
+        balancer._inflight_requests[pinned] += 10
+        moved = balancer.acquire_server("conv")[0]
+        assert moved != pinned
+        for _ in range(10):
+            balancer.release_server(pinned)
+        # Balanced again, so the current pin holds rather than bouncing back.
+        assert balancer.acquire_server("conv")[0] == moved
+
+    @pytest.mark.parametrize("bad", [-1, -0.5, float("-inf"), float("nan")])
+    def test_invalid_margins_are_rejected(self, bad):
+        # These break affinity even when the pin IS the least loaded, so affinity_broken
+        # would count moves that never happened.
+        with pytest.raises(ValueError, match="affinity_break_margin"):
+            lb(affinity_break_margin=bad)
+
+    def test_counters_cover_every_sticky_hit(self):
+        balancer = lb(affinity_break_margin=1)
+        balancer.acquire_server("conv")
+        for i in range(6):
+            balancer._inflight_requests["s0"] = i
+            balancer.acquire_server("conv")
+        status = balancer.get_status()
+        assert status["affinity_kept"] + status["affinity_broken"] == 6

verl/trainer/config/_generated_ppo_megatron_trainer.yaml

@@ -315,6 +315,7 @@ actor_rollout_ref:
     top_k: -1
     top_p: 1
     full_determinism: false
+    affinity_break_margin: null
     seed: 42
     prompt_length: ${oc.select:data.max_prompt_length,512}
     response_length: ${oc.select:data.max_response_length,512}

verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml

@@ -257,6 +257,7 @@ actor_rollout_ref:
     top_k: -1
     top_p: 1
     full_determinism: false
+    affinity_break_margin: null
     seed: 42
     prompt_length: ${oc.select:data.max_prompt_length,512}
     response_length: ${oc.select:data.max_response_length,512}

verl/trainer/config/_generated_ppo_trainer.yaml

@@ -268,6 +268,7 @@ actor_rollout_ref:
     top_k: -1
     top_p: 1
     full_determinism: false
+    affinity_break_margin: null
     seed: 42
     prompt_length: ${oc.select:data.max_prompt_length,512}
     response_length: ${oc.select:data.max_response_length,512}

verl/trainer/config/_generated_ppo_veomni_trainer.yaml

@@ -273,6 +273,7 @@ actor_rollout_ref:
     top_k: -1
     top_p: 1
     full_determinism: false
+    affinity_break_margin: null
     seed: 42
     prompt_length: ${oc.select:data.max_prompt_length,512}
     response_length: ${oc.select:data.max_response_length,512}

verl/trainer/config/rollout/rollout.yaml

@@ -26,6 +26,11 @@ top_p: 1
 # See https://pytorch.org/docs/stable/notes/randomness.html for details.
 full_determinism: false
 
+# Extra in-flight requests the replica holding a multi-turn session may have over the
+# least-loaded one before that session is re-routed. null never re-routes (unchanged
+# behavior); 0 always moves to the least-loaded replica. Ignored under full_determinism.
+affinity_break_margin: null
+
 # Random seed for rollout. Used as the seed for vLLM sampling and
 # enable_full_determinism() when full_determinism is True.
 seed: 42

verl/workers/config/rollout.py

@@ -170,6 +170,11 @@ class RolloutConfig(BaseConfig):
     # Whether to enable full determinism for reproducibility.
     full_determinism: bool = False
 
+    # Extra in-flight requests the replica holding a multi-turn session may have over the
+    # least-loaded one before that session is re-routed. null never re-routes (unchanged
+    # behavior), 0 always moves to the least-loaded replica. Ignored under full_determinism.
+    affinity_break_margin: Optional[int] = None
+
     # Random seed for rollout. Used as the seed for vLLM sampling and
     # enable_full_determinism() when full_determinism is True.
     seed: int = 42

verl/workers/rollout/llm_server.py

@@ -64,13 +64,17 @@ class GlobalRequestLoadBalancer:
       same request always routes to the same replica across runs.
     - **Dynamic Server Management**: Supports add/remove servers at runtime
       for hybrid scaling.
+    - **Bounded Affinity**: With ``affinity_break_margin`` set, a sticky session is
+      re-checked each turn and moved once its replica falls that many in-flight requests
+      behind the least-loaded one. Disabled (``None``) by default.
     """
 
     def __init__(
         self,
         servers: dict[str, ray.actor.ActorHandle],
         max_cache_size: int = DEFAULT_ROUTING_CACHE_SIZE,
         full_determinism: bool = False,
+        affinity_break_margin: Optional[int] = None,
     ):
         # Allow empty initial servers: in dynamic-resource-scheduling mode all
         # replicas are hybrid and will be registered later via add_servers().
@@ -79,10 +83,35 @@ def __init__(
         self._inflight_requests: dict[str, int] = {sid: 0 for sid in servers}
         self._request_id_to_server: LRUCache = LRUCache(maxsize=max_cache_size)
         self._full_determinism = full_determinism
+        # Written as `not (x >= 0)` so a NaN, which would break affinity even when the pin is
+        # least loaded, is rejected rather than silently accepted.
+        if affinity_break_margin is not None and not (affinity_break_margin >= 0):
+            raise ValueError(f"affinity_break_margin must be >= 0 or null, got {affinity_break_margin!r}")
+        self._affinity_break_margin = affinity_break_margin
+        self._affinity_kept = 0
+        self._affinity_broken = 0
+
+    def _least_loaded_server(self) -> str:
+        """Server with the fewest in-flight requests; ties go to the first registered."""
+        return min(self._inflight_requests, key=self._inflight_requests.get)
+
+    def _keep_affinity(self, server_id: str) -> bool:
+        """Whether ``server_id`` is still an acceptable home for its sticky session."""
+        # full_determinism routes by hash, so a load-dependent re-route would break it.
+        if self._affinity_break_margin is None or self._full_determinism:
+            return True
+        least = min(self._inflight_requests.values())
+        return self._inflight_requests[server_id] <= least + self._affinity_break_margin
 
     def acquire_server(self, request_id: str) -> tuple[str, ray.actor.ActorHandle]:
         """Acquire a server for the given request (sticky + least-loaded).
 
+        A sticky session keeps a conversation's prefix in one replica's KV cache. Held
+        unconditionally it also keeps feeding a replica that drew several long conversations
+        all of their later turns, since a cache hit returns before any load is consulted.
+        ``affinity_break_margin`` moves the session once its replica holds that many extra
+        in-flight requests: staying costs queueing delay, moving costs one re-prefill.
+
         Returns:
             A tuple of ``(server_id, actor_handle)`` in a single atomic call.
         """
@@ -91,6 +120,14 @@ def acquire_server(self, request_id: str) -> tuple[str, ray.actor.ActorHandle]:
             server_id = self._request_id_to_server[request_id]
             # Check if server is still in the active pool
             if server_id in self._inflight_requests:
+                if self._keep_affinity(server_id):
+                    self._affinity_kept += 1
+                    self._inflight_requests[server_id] += 1
+                    return server_id, self._servers[server_id]
+                self._affinity_broken += 1
+                # Re-pin, or the remaining turns re-prefill somewhere new every time.
+                server_id = self._least_loaded_server()
+                self._request_id_to_server[request_id] = server_id
                 self._inflight_requests[server_id] += 1
                 return server_id, self._servers[server_id]
             # Server was removed, clear stale cache entry and re-select
@@ -106,9 +143,7 @@ def acquire_server(self, request_id: str) -> tuple[str, ray.actor.ActorHandle]:
             # which varies run-to-run, so it is bypassed entirely here.
             server_id = list(self._servers)[hash(request_id) % len(self._servers)]
         else:
-            min_count = min(self._inflight_requests.values())
-            candidates = [sid for sid, count in self._inflight_requests.items() if count == min_count]
-            server_id = candidates[0]
+            server_id = self._least_loaded_server()
         self._request_id_to_server[request_id] = server_id
         self._inflight_requests[server_id] += 1
         return server_id, self._servers[server_id]
@@ -187,6 +222,10 @@ def get_status(self) -> dict:
             "total_inflight": sum(self._inflight_requests.values()),
             "active_servers": len(self._inflight_requests),
             "registered_handles": list(self._servers.keys()),
+            # Sticky hits only, not first placements.
+            "affinity_break_margin": self._affinity_break_margin,
+            "affinity_kept": self._affinity_kept,
+            "affinity_broken": self._affinity_broken,
         }
 
     def get_total_inflight(self) -> int:
@@ -603,11 +642,12 @@ async def _init_global_load_balancer(self) -> None:
             servers=dict(zip(self.server_addresses, self.server_handles, strict=True)),
             max_cache_size=DEFAULT_ROUTING_CACHE_SIZE,
         )
-        # The default GlobalRequestLoadBalancer honors the full_determinism flag
-        # in acquire_server. A custom subclass overrides acquire_server and takes
-        # full control of routing, so the flag is not forwarded to it.
+        # The default GlobalRequestLoadBalancer honors full_determinism and the affinity
+        # margin in acquire_server. A custom subclass overrides acquire_server and takes full
+        # control of routing, so neither is forwarded to it.
         if load_balancer_cls is GlobalRequestLoadBalancer:
             kwargs["full_determinism"] = getattr(self.rollout_config, "full_determinism", False)
+            kwargs["affinity_break_margin"] = getattr(self.rollout_config, "affinity_break_margin", None)
         self.global_load_balancer = ray.remote(load_balancer_cls).remote(**kwargs)
 
     def get_client(self, client_cls: type[LLMServerClient] | None = None, **kwargs) -> LLMServerClient: