[rollout, trainer, cfg] feat: per-request abort hooks and AbortableLLMServerClient

tests/workers/rollout/test_abortable_llm_client_on_cpu.py

@@ -0,0 +1,159 @@
+# Copyright 2024 Bytedance Ltd. 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.
+"""CPU unit tests for AbortableLLMServerClient: per-request abort + in-flight cleanup.
+
+These mock the Ray load balancer and server handles, so no GPU/vLLM is required.
+"""
+
+import asyncio
+from types import SimpleNamespace
+
+import pytest
+
+from verl.workers.rollout.llm_server import AbortableLLMServerClient
+from verl.workers.rollout.replica import TokenOutput
+
+# Generous guard against a hung task; not a performance assertion. The first asyncio test
+# in a heavily loaded CI run pays event-loop cold-start cost, so keep this well above 1s.
+WAIT_TIMEOUT = 30.0
+
+
+class _FakeRemoteMethod:
+    """Mimics a Ray actor method: ``handle.method.remote(...)``."""
+
+    def __init__(self, fn):
+        self._fn = fn
+
+    def remote(self, *args, **kwargs):
+        return self._fn(*args, **kwargs)
+
+
+class _FakeServer:
+    """Fake vLLMHttpServer handle.
+
+    ``generate`` blocks on ``release`` so the test can observe the in-flight state and
+    issue an abort while the request is mid-flight. ``abort_request`` records its arg.
+    """
+
+    def __init__(self, stop_reason="aborted"):
+        self.started = asyncio.Event()
+        self.release = asyncio.Event()
+        self.aborted = []
+        self._stop_reason = stop_reason
+
+        async def _generate(**kwargs):
+            self.started.set()
+            await self.release.wait()
+            return TokenOutput(token_ids=[1, 2, 3], stop_reason=self._stop_reason, extra_fields={})
+
+        def _abort(request_id):
+            # Record synchronously (like a Ray .remote() dispatch) so fire-and-forget callers
+            # that never await still register the abort. Return a resolved future so awaiting
+            # callers (abort()) work too.
+            self.aborted.append(request_id)
+            fut = asyncio.get_running_loop().create_future()
+            fut.set_result({"aborted": True})
+            return fut
+
+        self.generate = _FakeRemoteMethod(_generate)
+        self.abort_request = _FakeRemoteMethod(_abort)
+
+
+class _FakeLoadBalancer:
+    def __init__(self, server):
+        async def _acquire(request_id):
+            return ("server-0", server)
+
+        def _release(server_id):
+            # Fire-and-forget in the client; return a non-awaitable to avoid "coroutine
+            # never awaited" warnings.
+            return None
+
+        self.acquire_server = _FakeRemoteMethod(_acquire)
+        self.release_server = _FakeRemoteMethod(_release)
+
+
+def _make_client(server):
+    # config is only touched by generate() on the priority!=0 path, which we never hit.
+    return AbortableLLMServerClient(config=SimpleNamespace(), load_balancer_handle=_FakeLoadBalancer(server))
+
+
+@pytest.mark.asyncio
+async def test_record_abort_and_cleanup():
+    """Dispatch records the request, abort targets the inner id, completion clears it."""
+    server = _FakeServer(stop_reason="aborted")
+    client = _make_client(server)
+
+    task = asyncio.create_task(client.generate(request_id="req-1", prompt_ids=[1, 2], sampling_params={}))
+
+    # Wait until the request is actually dispatched to the (fake) server.
+    await asyncio.wait_for(server.started.wait(), timeout=WAIT_TIMEOUT)
+    assert "req-1" in client._inflight
+    inner_request_id, recorded_server = client._inflight["req-1"]
+    assert recorded_server is server
+
+    # Abort the specific request; it must hit the inner (vLLM) request id, not the outer one.
+    await client.abort("req-1")
+    assert server.aborted == [inner_request_id]
+    assert inner_request_id != "req-1"
+
+    # Entry is only cleared once generate() returns through its finally clause.
+    server.release.set()
+    output = await asyncio.wait_for(task, timeout=WAIT_TIMEOUT)
+    assert output.stop_reason == "aborted"
+    assert "req-1" not in client._inflight
+
+
+@pytest.mark.asyncio
+async def test_inflight_cleared_on_normal_completion():
+    """A request that finishes normally is also removed from the in-flight table."""
+    server = _FakeServer(stop_reason="stop")
+    client = _make_client(server)
+
+    task = asyncio.create_task(client.generate(request_id="req-2", prompt_ids=[1], sampling_params={}))
+    await asyncio.wait_for(server.started.wait(), timeout=WAIT_TIMEOUT)
+    assert "req-2" in client._inflight
+
+    server.release.set()
+    await asyncio.wait_for(task, timeout=WAIT_TIMEOUT)
+    assert "req-2" not in client._inflight
+
+
+@pytest.mark.asyncio
+async def test_abort_unknown_request_is_noop():
+    """Aborting an unknown/finished request neither raises nor calls the server."""
+    server = _FakeServer()
+    client = _make_client(server)
+
+    await client.abort("does-not-exist")
+    assert server.aborted == []
+
+
+@pytest.mark.asyncio
+async def test_cancellation_aborts_server_and_clears_inflight():
+    """Cancelling generate() (e.g. timeout) aborts the server-side request, no leak."""
+    server = _FakeServer()
+    client = _make_client(server)
+
+    task = asyncio.create_task(client.generate(request_id="req-3", prompt_ids=[1], sampling_params={}))
+    await asyncio.wait_for(server.started.wait(), timeout=WAIT_TIMEOUT)
+    inner_request_id, _ = client._inflight["req-3"]
+
+    task.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await task
+
+    # _on_cancel fired the abort to the server, and _on_complete cleared the entry.
+    assert server.aborted == [inner_request_id]
+    assert "req-3" not in client._inflight

verl/experimental/agent_loop/agent_loop.py

@@ -1212,11 +1212,11 @@ async def generate_sequences(self, prompts: DataProto) -> DataProto:
         if "priority" not in prompts.non_tensor_batch:
             prompts.non_tensor_batch["priority"] = np.arange(len(prompts), dtype=np.int64)
 
-        chunkes = prompts.chunk(len(self.agent_loop_workers))
+        chunks = prompts.chunk(len(self.agent_loop_workers))
         outputs = await asyncio.gather(
             *[
                 worker.generate_sequences.remote(chunk)
-                for worker, chunk in zip(self.agent_loop_workers, chunkes, strict=True)
+                for worker, chunk in zip(self.agent_loop_workers, chunks, strict=True)
             ]
         )
         output = DataProto.concat(outputs)

verl/trainer/ppo/v1/trainer_base.py

@@ -76,7 +76,7 @@
 from verl.utils.debug import marked_timer
 from verl.utils.debug.metrics import calculate_debug_metrics
 from verl.utils.fs import copy_to_local
-from verl.utils.import_utils import load_extern_type
+from verl.utils.import_utils import load_class_from_fqn, load_extern_type
 from verl.utils.metric import reduce_metrics
 from verl.utils.py_functional import rename_dict
 from verl.utils.seqlen_balancing import calculate_workload, get_seqlen_balanced_partitions, log_seqlen_unbalance
@@ -284,7 +284,12 @@ def _setup(self):
 
     def get_llm_client(self) -> LLMServerClient:
         """Get the LLM server client for rollout generation."""
-        return self.llm_server_manager.get_client()
+        # Support custom LLMServerClient via config (mirrors agent_loop_manager_class).
+        llm_client_class_fqn = self.config.actor_rollout_ref.rollout.get("agent", {}).get("llm_client_class")
+        llm_client_kwargs = (
+            {"client_cls": load_class_from_fqn(llm_client_class_fqn, "LLMServerClient")} if llm_client_class_fqn else {}
+        )
+        return self.llm_server_manager.get_client(**llm_client_kwargs)
 
     def get_teacher_client(self) -> Optional[dict[str, LLMServerClient]]:
         """Get the On-Policy Distillation teacher server clients.

verl/workers/config/rollout.py

@@ -77,6 +77,9 @@ class AgentLoopConfig(BaseConfig):
     # Fully qualified class name for custom AgentLoopManager (e.g., "mypackage.module.MyManager").
     # Security: This class will be dynamically imported via importlib. Only use trusted class paths.
     agent_loop_manager_class: Optional[str] = None
+    # Fully qualified class name for custom LLMServerClient (e.g., "mypackage.module.MyClient").
+    # Security: This class will be dynamically imported via importlib. Only use trusted class paths.
+    llm_client_class: Optional[str] = None
 
 
 @dataclass

verl/workers/rollout/llm_server.py

@@ -220,6 +220,8 @@ async def generate(
             TokenOutput | DiffusionOutput: token or diffusion output
         """
         server_id, server = await self._acquire_server(request_id)
+        inner_request_id = uuid4().hex  # hoisted from the .remote() call so subclasses can record it
+        self._on_dispatch(request_id, inner_request_id, server)
         try:
             multimodal_kwargs = {}
             if audio_data is not None:
@@ -232,7 +234,7 @@ async def generate(
                 {"priority": priority} if priority != 0 and self.config.actor_rollout_ref.rollout.name == "vllm" else {}
             )
             output: TokenOutput = await server.generate.remote(
-                request_id=uuid4().hex,  # use new request_id for each turn
+                request_id=inner_request_id,
                 prompt_ids=prompt_ids,
                 sampling_params=sampling_params,
                 image_data=image_data,
@@ -245,9 +247,27 @@ async def generate(
             output.extra_fields.setdefault("min_global_steps", global_steps)
             output.extra_fields.setdefault("max_global_steps", global_steps)
             return output
+        except asyncio.CancelledError:
+            # Client-side cancellation (e.g. asyncio.wait_for/timeout) does not reach the
+            # server, so the server-side request would keep consuming compute. Fire this
+            # hook while the in-flight entry still exists (before the finally clause runs
+            # _on_complete) so subclasses can abort it on the server.
+            self._on_cancel(request_id)
+            raise
         finally:
+            self._on_complete(request_id)
             self._release_server(server_id)
 
+    def _on_dispatch(self, request_id, inner_request_id, server):
+        """Hook fired after the inner vLLM request_id is assigned, before awaiting generation.
+        Default no-op; subclasses may record (request_id/inner_request_id, server) to enable abort."""
+
+    def _on_complete(self, request_id):
+        """Hook fired when generation finishes or raises. Default no-op."""
+
+    def _on_cancel(self, request_id):
+        """Hook fired when generation is cancelled (CancelledError). Default no-op."""
+
 
 class FullyAsyncLLMServerClient(LLMServerClient):
     """FullyLLMServerClient supports resume generation on partial rollout, making rollout interruption
@@ -353,6 +373,62 @@ async def generate(
         return final_output
 
 
+class AbortableLLMServerClient(LLMServerClient):
+    """LLMServerClient that tracks in-flight requests so individual ones can be aborted.
+
+    Enables selective (per-request) abort, e.g. cancelling a single timed-out rollout or
+    letting a custom AgentLoopWorker/AgentLoopManager interrupt a specific request. This is
+    complementary to the broadcast ``abort_all_requests`` on the rollout replica.
+
+    Inherits the base (non-resuming) ``generate``: once a request is aborted it returns with
+    ``stop_reason="aborted"`` rather than resuming, which matches hard-abort semantics.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # request_id -> (inner_request_id, server handle). asyncio is single-threaded,
+        # so a plain dict needs no locking.
+        self._inflight: dict[str, tuple[str, ray.actor.ActorHandle]] = {}
+
+    def _on_dispatch(self, request_id, inner_request_id, server):
+        self._inflight[request_id] = (inner_request_id, server)
+
+    def _on_complete(self, request_id):
+        # Sole owner of removal: fires from generate()'s finally on success, abort, or error.
+        self._inflight.pop(request_id, None)
+
+    def _on_cancel(self, request_id):
+        # Cancellation (e.g. timeout) never reaches the server on its own; fire-and-forget an
+        # abort so the server-side request stops. We cannot safely await during cancellation,
+        # and removal is still handled by _on_complete in the finally clause.
+        self._send_abort(request_id)
+
+    def _send_abort(self, request_id: str):
+        """Dispatch an abort RPC for ``request_id`` to the owning server (does not await).
+
+        Returns the abort ObjectRef, or ``None`` if the request is unknown (already finished,
+        aborted, or not yet dispatched). Does not remove the in-flight entry; ``_on_complete``
+        owns removal.
+        """
+        entry = self._inflight.get(request_id)
+        if entry is None:
+            return None
+        inner_request_id, server = entry
+        return server.abort_request.remote(inner_request_id)
+
+    async def abort(self, request_id: str) -> None:
+        """Abort a single in-flight request by its (outer) request_id and await the result.
+
+        No-op if the request is unknown. The in-flight entry is removed by ``_on_complete``
+        when the aborted ``generate`` returns through its finally clause. Client-side
+        cancellation of ``generate`` (e.g. ``asyncio.wait_for``) is handled automatically via
+        ``_on_cancel``, so callers do not need to abort before cancelling.
+        """
+        ref = self._send_abort(request_id)
+        if ref is not None:
+            await ref
+
+
 class LLMServerManager:
     """LLMServerManager is responsible for:
     - Launch server replicas