verl-project
diff --git a/‎tests/checkpoint_engine/test_special_server_adapter.py‎
Lines changed: 4 additions & 3 deletions b/‎tests/checkpoint_engine/test_special_server_adapter.py‎
Lines changed: 4 additions & 3 deletions
diff --git a/‎tests/checkpoint_engine/test_utils.py‎
Lines changed: 1 addition & 1 deletion b/‎tests/checkpoint_engine/test_utils.py‎
Lines changed: 1 addition & 1 deletion
diff --git a/‎tests/experimental/agent_loop/test_basic_agent_loop.py‎
Lines changed: 105 additions & 20 deletions b/‎tests/experimental/agent_loop/test_basic_agent_loop.py‎
Lines changed: 105 additions & 20 deletions
diff --git a/‎verl/checkpoint_engine/base.py‎
Lines changed: 39 additions & 8 deletions b/‎verl/checkpoint_engine/base.py‎
Lines changed: 39 additions & 8 deletions
diff --git a/‎verl/experimental/fully_async_policy/detach_utils.py‎
Lines changed: 1 addition & 10 deletions b/‎verl/experimental/fully_async_policy/detach_utils.py‎
Lines changed: 1 addition & 10 deletions
@@ -21,12 +21,13 @@
 
 from tests.checkpoint_engine.test_utils import create_trainer_worker_group
 from verl.checkpoint_engine import CheckpointEngineManager
+from verl.experimental.fully_async_policy.fully_async_rollouter import FullyAsyncLLMServerClient
 from verl.single_controller.ray import (
     RayResourcePool,
 )
 from verl.utils.config import omega_conf_to_dataclass
 from verl.workers.config import CheckpointEngineConfig, HFModelConfig
-from verl.workers.rollout.llm_server import FullyLLMServerClient, LLMServerClient, LLMServerManager
+from verl.workers.rollout.llm_server import LLMServerClient, LLMServerManager
 
 
 @pytest.fixture
@@ -123,7 +124,7 @@ async def _run_server_manager_without_resume(
 async def _run_server_manager_with_resume(
     initial_steps: int,
     train_steps: int,
-    server_manager: FullyLLMServerClient,
+    server_manager: FullyAsyncLLMServerClient,
     checkpoint_manager: CheckpointEngineManager,
     prompts: list[list[dict]],
     tokenizer: PreTrainedTokenizer,
@@ -231,7 +232,7 @@ async def test_server_adapter(init_config):
     await _run_server_manager_with_resume(
         initial_steps=4,
         train_steps=3,
-        server_manager=llm_server_manager.get_client(fully_async=True),
+        server_manager=llm_server_manager.get_client(client_cls=FullyAsyncLLMServerClient),
         checkpoint_manager=checkpoint_manager,
         prompts=prompts,
         tokenizer=model_config.tokenizer,
 
@@ -40,7 +40,7 @@ def __init__(self, config: TrainingWorkerConfig, checkpoint_engine_config: Check
         self.checkpoint_engine = CheckpointEngineRegistry.new(backend, bucket_size=bucket_size, **engine_kwargs)
 
     @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False)
-    async def update_weights(self, global_steps: int = None):
+    async def update_weights(self, global_steps: int = None, mode: str = "auto"):
         per_tensor_param, _ = self.engine.get_per_tensor_param()
         await self.checkpoint_engine.send_weights(per_tensor_param)
 
 
@@ -422,51 +422,136 @@ class TestLoadBalancerRouting:
 
     def test_distributes_across_servers(self, ray_for_lb):
         lb = GlobalRequestLoadBalancer.remote(servers={"s0": None, "s1": None, "s2": None})
-        servers = [ray.get(lb.acquire_server.remote(request_id=f"r{i}")) for i in range(3)]
+        servers = [ray.get(lb.acquire_server.remote(request_id=f"r{i}"))[0] for i in range(3)]
         assert sorted(servers) == ["s0", "s1", "s2"]
 
     def test_new_requests_route_to_least_loaded(self, ray_for_lb):
         lb = GlobalRequestLoadBalancer.remote(servers={"s0": None, "s1": None, "s2": None})
         # Load s0 with 3 inflight requests
-        ray.get(lb.acquire_server.remote(request_id="a"))  # -> s0
-        ray.get(lb.acquire_server.remote(request_id="a"))  # sticky -> s0
-        ray.get(lb.acquire_server.remote(request_id="a"))  # sticky -> s0
+        ray.get(lb.acquire_server.remote(request_id="a"))[0]  # -> s0
+        ray.get(lb.acquire_server.remote(request_id="a"))[0]  # sticky -> s0
+        ray.get(lb.acquire_server.remote(request_id="a"))[0]  # sticky -> s0
         # Load s1 with 1 inflight request
-        ray.get(lb.acquire_server.remote(request_id="b"))  # -> s1
+        ray.get(lb.acquire_server.remote(request_id="b"))[0]  # -> s1
         # s2 has 0 inflight, so next new request must go to s2
-        s_new = ray.get(lb.acquire_server.remote(request_id="d"))
+        s_new = ray.get(lb.acquire_server.remote(request_id="d"))[0]
         assert s_new == "s2"
 
     def test_release_rebalances(self, ray_for_lb):
         lb = GlobalRequestLoadBalancer.remote(servers={"s0": None, "s1": None})
-        s0 = ray.get(lb.acquire_server.remote(request_id="r0"))
-        s1 = ray.get(lb.acquire_server.remote(request_id="r1"))
+        s0 = ray.get(lb.acquire_server.remote(request_id="r0"))[0]
+        s1 = ray.get(lb.acquire_server.remote(request_id="r1"))[0]
         assert s0 != s1
         ray.get(lb.release_server.remote(server_id=s0))
         ray.get(lb.release_server.remote(server_id=s1))
-        s2 = ray.get(lb.acquire_server.remote(request_id="r2"))
-        s3 = ray.get(lb.acquire_server.remote(request_id="r3"))
+        s2 = ray.get(lb.acquire_server.remote(request_id="r2"))[0]
+        s3 = ray.get(lb.acquire_server.remote(request_id="r3"))[0]
         assert s2 != s3
 
-    def test_release_invalid_server_raises(self, ray_for_lb):
+    def test_release_invalid_server_silently_ignored(self, ray_for_lb):
+        """Releasing a nonexistent server is silently ignored (hybrid-safe)."""
         lb = GlobalRequestLoadBalancer.remote(servers={"s0": None, "s1": None})
-        with pytest.raises(ray.exceptions.RayTaskError, match="Invalid server_id") as excinfo:
-            ray.get(lb.release_server.remote(server_id="nonexistent"))
-        assert "Invalid server_id" in str(excinfo.value)
+        # Should not raise
+        ray.get(lb.release_server.remote(server_id="nonexistent"))
 
-    def test_release_without_inflight_raises(self, ray_for_lb):
+    def test_release_without_inflight_silently_ignored(self, ray_for_lb):
+        """Releasing a server with no inflight requests is silently ignored (hybrid-safe)."""
         lb = GlobalRequestLoadBalancer.remote(servers={"s0": None, "s1": None})
-        with pytest.raises(ray.exceptions.RayTaskError, match="no inflight") as excinfo:
-            ray.get(lb.release_server.remote(server_id="s1"))
-        assert "no inflight" in str(excinfo.value)
+        # Should not raise even though s1 has 0 inflight
+        ray.get(lb.release_server.remote(server_id="s1"))
 
 
 class TestLoadBalancerStickySession:
     """Request-level sticky session."""
 
     def test_same_request_id_same_server(self, ray_for_lb):
         lb = GlobalRequestLoadBalancer.remote(servers={"s0": None, "s1": None, "s2": None, "s3": None})
-        s0 = ray.get(lb.acquire_server.remote(request_id="conv-abc"))
+        s0 = ray.get(lb.acquire_server.remote(request_id="conv-abc"))[0]
         ray.get(lb.release_server.remote(server_id=s0))
-        s1 = ray.get(lb.acquire_server.remote(request_id="conv-abc"))
+        s1 = ray.get(lb.acquire_server.remote(request_id="conv-abc"))[0]
         assert s0 == s1
+
+
+class TestLoadBalancerHybrid:
+    """Dynamic server add/remove for hybrid scaling."""
+
+    def test_add_server(self, ray_for_lb):
+        lb = GlobalRequestLoadBalancer.remote(servers={"s0": None, "s1": None})
+        ray.get(lb.add_servers.remote(servers={"s2": None}))
+        status = ray.get(lb.get_status.remote())
+        assert "s2" in status["servers"]
+        assert status["servers"]["s2"] == 0
+
+    def test_remove_server_purges_handle(self, ray_for_lb):
+        lb = GlobalRequestLoadBalancer.remote(servers={"s0": None, "s1": None})
+        ray.get(lb.remove_servers.remote(server_ids=["s1"]))
+        # remove_server now purges from both _inflight_requests and _servers
+        status = ray.get(lb.get_status.remote())
+        assert "s1" not in status["servers"]
+        assert "s1" not in status["registered_handles"]
+        # New requests should only go to s0
+        s = ray.get(lb.acquire_server.remote(request_id="r1"))[0]
+        assert s == "s0"
+
+    def test_removed_server_invalidates_sticky_session(self, ray_for_lb):
+        """When a sticky session points to a removed server, cache is invalidated."""
+        lb = GlobalRequestLoadBalancer.remote(servers={"s0": None, "s1": None})
+        # Occupy s0 so that the sticky request is assigned to s1
+        ray.get(lb.acquire_server.remote(request_id="occupy-s0"))[0]  # -> s0
+        # Pin request to s1 (least-loaded now)
+        s1 = ray.get(lb.acquire_server.remote(request_id="sticky-req"))[0]
+        assert s1 == "s1"
+        ray.get(lb.release_server.remote(server_id=s1))
+        # Remove s1
+        ray.get(lb.remove_servers.remote(server_ids=["s1"]))
+        # Sticky session should be invalidated and reroute to s0
+        s_new = ray.get(lb.acquire_server.remote(request_id="sticky-req"))[0]
+        assert s_new == "s0"
+
+    def test_remove_server_also_purges_registry(self, ray_for_lb):
+        """remove_servers atomically purges from both LB pool and handle registry."""
+        lb = GlobalRequestLoadBalancer.remote(servers={"s0": None, "s1": None})
+        ray.get(lb.remove_servers.remote(server_ids=["s1"]))
+        status = ray.get(lb.get_status.remote())
+        # Both _inflight_requests and _servers are cleaned up (no separate cleanup step needed)
+        assert "s1" not in status["servers"]
+        assert "s1" not in status["registered_handles"]
+
+    def test_get_all_servers_excludes_removed(self, ray_for_lb):
+        lb = GlobalRequestLoadBalancer.remote(servers={"s0": None, "s1": None, "s2": None})
+        ray.get(lb.remove_servers.remote(server_ids=["s1"]))
+        all_servers = ray.get(lb.get_all_servers.remote())
+        assert "s0" in all_servers
+        assert "s2" in all_servers
+        assert "s1" not in all_servers
+
+    def test_no_available_servers_raises(self, ray_for_lb):
+        lb = GlobalRequestLoadBalancer.remote(servers={"s0": None, "s1": None})
+        ray.get(lb.remove_servers.remote(server_ids=["s0", "s1"]))
+        with pytest.raises(ray.exceptions.RayTaskError, match="No available servers"):
+            ray.get(lb.acquire_server.remote(request_id="r1"))
+
+    def test_add_server_readds_previously_removed(self, ray_for_lb):
+        """Re-adding a previously removed server makes it routable again."""
+        lb = GlobalRequestLoadBalancer.remote(servers={"s0": None, "s1": None})
+        ray.get(lb.remove_servers.remote(server_ids=["s1"]))
+        # s1 is removed, only s0 is available
+        assert ray.get(lb.acquire_server.remote(request_id="r1"))[0] == "s0"
+        # Re-add s1
+        ray.get(lb.add_servers.remote(servers={"s1": None}))
+        # Now both s0 and s1 should be available
+        s = ray.get(lb.acquire_server.remote(request_id="r2"))[0]
+        assert s in ("s0", "s1")
+
+    def test_get_inflight_count(self, ray_for_lb):
+        lb = GlobalRequestLoadBalancer.remote(servers={"s0": None, "s1": None})
+        assert ray.get(lb.get_inflight_count.remote(server_id="s0")) == 0
+        ray.get(lb.acquire_server.remote(request_id="r1"))[0]  # -> s0 (least loaded)
+        assert ray.get(lb.get_inflight_count.remote(server_id="s0")) == 1
+
+    def test_get_status_reports_active_correctly(self, ray_for_lb):
+        lb = GlobalRequestLoadBalancer.remote(servers={"s0": None, "s1": None, "s2": None})
+        ray.get(lb.remove_servers.remote(server_ids=["s1"]))
+        status = ray.get(lb.get_status.remote())
+        assert status["active_servers"] == 2  # s0 and s2
+        assert status["total_inflight"] == 0
@@ -416,6 +416,34 @@ async def wake_up_replicas(self):
         """Resume all rollout replicas: recover kv_cache and weights device memory."""
         await asyncio.gather(*[r.wake_up() for r in self.replicas])
 
+    @auto_await
+    async def abort_replicas(self):
+        """Abort all in-flight requests on every replica."""
+        await asyncio.gather(*[r.abort_all_requests() for r in self.replicas])
+
+    @auto_await
+    async def resume_generation_replicas(self):
+        """Resume generation on all replicas after abort_all_requests."""
+        await asyncio.gather(*[r.resume_generation() for r in self.replicas])
+
+    @auto_await
+    async def release_kv_cache_replicas(self):
+        """Release kv_cache of all rollout replicas before NCCL weight sync.
+
+        Unlike sleep_replicas(), this only frees the kv_cache and leaves model
+        weights untouched, so the NCCL transfer can write directly into the
+        existing weight buffers.  Call resume_kv_cache_replicas() after sync.
+        """
+        await asyncio.gather(*[r.release_kv_cache() for r in self.replicas])
+
+    @auto_await
+    async def resume_kv_cache_replicas(self):
+        """Restore kv_cache of all rollout replicas after NCCL weight sync.
+
+        Counterpart to release_kv_cache_replicas().
+        """
+        await asyncio.gather(*[r.resume_kv_cache() for r in self.replicas])
+
     @auto_await
     async def update_weights(self, global_steps: int = None):
         """Update weights from trainer to rollout replicas.
@@ -426,11 +454,11 @@ async def update_weights(self, global_steps: int = None):
 
         # 0. update weights for sync training with colocated trainer and rollout
         if self.backend == "naive":
-            ray.get(self.trainer.update_weights(global_steps=global_steps))
+            ray.get(self.trainer.update_weights(global_steps=global_steps, mode=self.backend))
             return
 
         # 1. abort and save all unfinished requests for partial rollout
-        await asyncio.gather(*[r.abort_all_requests() for r in self.replicas])
+        await self.abort_replicas()
 
         # 2. create a temporay worker group for all replicas
         workers = []
@@ -439,26 +467,29 @@ async def update_weights(self, global_steps: int = None):
         rollout = RayWorkerGroup(worker_handles=workers, ray_cls_with_init=RayClassWithInitArgs(cls=_worker_cls))
         trainer = self.trainer
 
-        # 3. sleep replicas to free kv_cache before weight sync (if free_cache_engine is enabled)
-        await self.sleep_replicas()
+        # 3. release kv_cache before weight sync (weights stay in place)
+        await self.release_kv_cache_replicas()
 
         # 4. build process group
         self.build_process_group(rollout)
 
         # 5. update weights of all workers
-        ray.get(trainer.update_weights(global_steps=global_steps) + rollout.update_weights(global_steps=global_steps))
+        ray.get(
+            trainer.update_weights(global_steps=global_steps, mode=self.backend)
+            + rollout.update_weights(global_steps=global_steps)
+        )
 
         # 6. finalize all workers
         ray.get(
             trainer.execute_checkpoint_engine(["finalize"] * trainer.world_size)
             + rollout.execute_checkpoint_engine(["finalize"] * rollout.world_size)
         )
 
-        # 7. resume replicas to recover kv_cache (for free_cache_engine scenarios)
-        await self.wake_up_replicas()
+        # 7. restore kv_cache after weight sync
+        await self.resume_kv_cache_replicas()
 
         # 8. resume all unfinished requests for partial rollout
-        await asyncio.gather(*[r.resume_generation() for r in self.replicas])
+        await self.resume_generation_replicas()
 
 
 async def split_weight_chunks(
 
@@ -15,7 +15,7 @@
 import time
 from collections import defaultdict
 from dataclasses import dataclass
-from typing import Any, Optional
+from typing import Any
 
 import numpy as np
 import torch
@@ -39,15 +39,6 @@ class RolloutSample:
     rollout_status: dict[str, Any]
 
 
-@dataclass
-class ValidateMetrics:
-    """Metrics for validation"""
-
-    timing_raw: dict[str, Any]
-    metrics: Optional[dict[str, Any]] = None
-    val_generations: Optional[list[tuple]] = None
-
-
 def prepare_single_generation_data(batch_dict, config) -> DataProto:
     """
     Similar to the logic of ray_trainer._prepare_generate_batch, but for a single sample.