diff --git a/tests/workers/rollout/test_abortable_llm_client_on_cpu.py b/tests/workers/rollout/test_abortable_llm_client_on_cpu.py new file mode 100644 index 00000000000..1d662eaa4d0 --- /dev/null +++ b/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 diff --git a/verl/experimental/agent_loop/agent_loop.py b/verl/experimental/agent_loop/agent_loop.py index 9b618f89d2f..55054028523 100644 --- a/verl/experimental/agent_loop/agent_loop.py +++ b/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) diff --git a/verl/trainer/ppo/v1/trainer_base.py b/verl/trainer/ppo/v1/trainer_base.py index b40c4853a94..8348165c17e 100644 --- a/verl/trainer/ppo/v1/trainer_base.py +++ b/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. diff --git a/verl/workers/config/rollout.py b/verl/workers/config/rollout.py index 74b2f755a22..137713ad807 100644 --- a/verl/workers/config/rollout.py +++ b/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 diff --git a/verl/workers/rollout/llm_server.py b/verl/workers/rollout/llm_server.py index e85d3029ad1..c0b6c36e3f9 100644 --- a/verl/workers/rollout/llm_server.py +++ b/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