From e4290f4715268faf471ad301a43405336b72c3d0 Mon Sep 17 00:00:00 2001 From: "tongyuxuan.361" Date: Wed, 14 Oct 2026 19:22:39 +0800 Subject: [PATCH] [framework, gateway] fix: settle cancelled session creation without losing its owner --- .../test_generate_sequences_on_cpu.py | 32 +++ .../gateway/test_gateway_manager_on_cpu.py | 198 ++++++++++++++++++ uni_agent/framework/framework.py | 5 +- uni_agent/gateway/manager.py | 48 ++++- 4 files changed, 274 insertions(+), 9 deletions(-) diff --git a/tests/uni_agent/framework/test_generate_sequences_on_cpu.py b/tests/uni_agent/framework/test_generate_sequences_on_cpu.py index 8ec4a77a6..a15eec537 100644 --- a/tests/uni_agent/framework/test_generate_sequences_on_cpu.py +++ b/tests/uni_agent/framework/test_generate_sequences_on_cpu.py @@ -2126,3 +2126,35 @@ async def scored_runner(**kwargs): }, ] assert "runner_reward_info" not in framework._trajectory_meta(trajectories[0]) + + +@pytest.mark.cpu +@pytest.mark.level0 +@pytest.mark.asyncio +async def test_runner_error_survives_gateway_abort_failure(): + async def failing_runner(**kwargs): + raise ValueError("runner exploded") + + class _GatewayManager(_FakeGatewayManager): + async def abort_session(self, session_id: str) -> None: + await super().abort_session(session_id) + raise RuntimeError("abort failed") + + runtime = _GatewayManager({"session-sample-0-rollout-0": [_trajectory()]}) + framework = await _build_framework_with_agent_runners( + agent_runners={"runner": _inline_runner_config(failing_runner)}, + gateway_manager=runtime, + ) + + # The abort failure is logged; the runner error is what callers see. + with pytest.raises(ValueError, match="runner exploded"): + await framework._run_agent_episode( + sample_fields={"raw_prompt": [], "uid": "uid-0"}, + sample_index=0, + session_index=0, + global_steps=7, + runner_name="runner", + runner_config=framework.runner_registry["runner"], + sampling_params={}, + ) + assert len(runtime.aborted_sessions) == 1 diff --git a/tests/uni_agent/gateway/test_gateway_manager_on_cpu.py b/tests/uni_agent/gateway/test_gateway_manager_on_cpu.py index 8e31c30e6..996f86592 100644 --- a/tests/uni_agent/gateway/test_gateway_manager_on_cpu.py +++ b/tests/uni_agent/gateway/test_gateway_manager_on_cpu.py @@ -325,3 +325,201 @@ async def generate( ] finally: await manager.shutdown() + + +class _LifecycleGateway: + """Controllable remote-method stub for close failures.""" + + def __init__(self, *, finalize_error=None, abort_error=None): + self.finalize_error = finalize_error + self.abort_error = abort_error + self.abort_calls = [] + self.finalize_calls = [] + + async def _abort(self, session_id): + self.abort_calls.append(session_id) + if self.abort_error is not None: + raise self.abort_error + + async def _finalize(self, session_id): + self.finalize_calls.append(session_id) + if self.finalize_error is not None: + raise self.finalize_error + return ["trajectory"] + + abort_session = property(lambda self: _FakeRemoteMethod(self._abort)) + finalize_session = property(lambda self: _FakeRemoteMethod(self._finalize)) + + +def _new_manager(gateway, *, session_ids=()): + from uni_agent.gateway.manager import GatewayManager + + manager = GatewayManager.__new__(GatewayManager) + manager.gateways = [gateway] + manager.active_sessions_per_gateway = [len(session_ids)] + manager._session_to_gateway_index = dict.fromkeys(session_ids, 0) + return manager + + +class _DetachedRemoteMethod: + """Ray-like call: cancelling the caller's await does not stop the actor method.""" + + def __init__(self, fn): + self._fn = fn + + def remote(self, *args, **kwargs): + return asyncio.shield(asyncio.ensure_future(self._fn(*args, **kwargs))) + + +class _SlowCreateGateway(_LifecycleGateway): + def __init__(self, *, create_error=None): + super().__init__() + self.create_error = create_error + self.created = [] + self.create_started = asyncio.Event() + self.create_finished = asyncio.Event() + self.release_create = asyncio.Event() + + async def _create(self, session_id, **kwargs): + self.create_started.set() + await self.release_create.wait() + self.create_finished.set() + if self.create_error is not None: + raise self.create_error + self.created.append(session_id) + return session_id + + create_session = property(lambda self: _DetachedRemoteMethod(self._create)) + abort_session = property(lambda self: _DetachedRemoteMethod(self._abort)) + + +@pytest.mark.cpu +@pytest.mark.level0 +@pytest.mark.asyncio +@pytest.mark.parametrize("create_error", [None, RuntimeError("create failed")]) +async def test_gateway_manager_cancelled_create_aborts_what_the_actor_created(create_error): + """After the caller cancels, only the manager can still abort the actor-side session.""" + gateway = _SlowCreateGateway(create_error=create_error) + manager = _new_manager(gateway) + + create = asyncio.create_task(manager.create_session("session-cancelled")) + await gateway.create_started.wait() + create.cancel() + await asyncio.sleep(0) + gateway.release_create.set() + with pytest.raises(asyncio.CancelledError): + await create + await asyncio.wait_for(gateway.create_finished.wait(), timeout=1) + + assert gateway.abort_calls == ([] if create_error else ["session-cancelled"]) + assert manager._session_to_gateway_index == {} + assert manager.active_sessions_per_gateway == [0] + + +@pytest.mark.cpu +@pytest.mark.level0 +@pytest.mark.asyncio +async def test_gateway_manager_cancelled_create_retains_owner_after_abort_failure(): + gateway = _SlowCreateGateway() + gateway.abort_error = RuntimeError("session abort failed before removal") + manager = _new_manager(gateway) + create = asyncio.create_task(manager.create_session("session-cancelled")) + await gateway.create_started.wait() + create.cancel() + gateway.release_create.set() + with pytest.raises(asyncio.CancelledError): + await create + + assert manager._session_to_gateway_index == {"session-cancelled": 0} + assert manager.active_sessions_per_gateway == [1] + gateway.abort_error = None + await manager.abort_session("session-cancelled") + assert gateway.abort_calls == ["session-cancelled"] * 2 + assert manager.active_sessions_per_gateway == [0] + + +@pytest.mark.cpu +@pytest.mark.level0 +@pytest.mark.asyncio +async def test_gateway_manager_repeated_cancellation_still_settles_remote_create(): + gateway = _SlowCreateGateway() + manager = _new_manager(gateway) + create = asyncio.create_task(manager.create_session("session-cancelled")) + await gateway.create_started.wait() + for _ in range(3): + create.cancel() + await asyncio.sleep(0) + still_waiting = not create.done() + gateway.release_create.set() + with pytest.raises(asyncio.CancelledError): + await create + await asyncio.wait_for(gateway.create_finished.wait(), timeout=1) + + assert still_waiting + assert gateway.abort_calls == ["session-cancelled"] + assert manager._session_to_gateway_index == {} + assert manager.active_sessions_per_gateway == [0] + + +@pytest.mark.cpu +@pytest.mark.level0 +@pytest.mark.asyncio +async def test_gateway_manager_close_failure_retains_mapping_and_owner(): + close_error = RuntimeError("actor-close-unknown") + gateway = _LifecycleGateway(finalize_error=close_error) + manager = _new_manager(gateway, session_ids=("session-close",)) + + with pytest.raises(RuntimeError, match="actor-close-unknown") as raised: + await manager.finalize_session("session-close") + + assert raised.value is close_error + assert gateway.finalize_calls == ["session-close"] + assert manager._session_to_gateway_index == {"session-close": 0} + assert manager.active_sessions_per_gateway == [1] + + await manager.abort_session("session-close") + + assert gateway.abort_calls == ["session-close"] + assert manager._session_to_gateway_index == {} + assert manager.active_sessions_per_gateway == [0] + + +@pytest.mark.cpu +@pytest.mark.level0 +@pytest.mark.asyncio +async def test_gateway_manager_rejects_duplicate_create_without_touching_owner(): + gateway = _SlowCreateGateway() + manager = _new_manager(gateway) + create = asyncio.create_task(manager.create_session("same-id")) + await gateway.create_started.wait() + + with pytest.raises(RuntimeError, match="already exists"): + await manager.create_session("same-id") + assert manager._session_to_gateway_index == {"same-id": 0} + assert manager.active_sessions_per_gateway == [1] + assert gateway.abort_calls == [] + + gateway.release_create.set() + assert await create == "same-id" + await manager.abort_session("same-id") + assert manager.active_sessions_per_gateway == [0] + + +@pytest.mark.cpu +@pytest.mark.level0 +@pytest.mark.asyncio +async def test_gateway_manager_concurrent_closes_release_load_slot_once(): + class _ConcurrentCloseGateway(_LifecycleGateway): + async def _finalize(self, session_id): + await asyncio.sleep(0) + return await super()._finalize(session_id) + + async def _abort(self, session_id): + await asyncio.sleep(0) + await super()._abort(session_id) + + gateway = _ConcurrentCloseGateway() + manager = _new_manager(gateway, session_ids=("session-close",)) + await asyncio.gather(manager.finalize_session("session-close"), manager.abort_session("session-close")) + assert manager._session_to_gateway_index == {} + assert manager.active_sessions_per_gateway == [0] diff --git a/uni_agent/framework/framework.py b/uni_agent/framework/framework.py index 7ec064e02..09e9f29ae 100644 --- a/uni_agent/framework/framework.py +++ b/uni_agent/framework/framework.py @@ -869,7 +869,10 @@ async def _run_agent_episode( raise except Exception: logger.exception("session %s failed (runner=%s); aborting session", session_id, runner_name) - await self.gateway_manager.abort_session(session_id) + try: + await self.gateway_manager.abort_session(session_id) + except Exception: + logger.exception("session %s: Gateway abort failed after runner failure", session_id) session_trace.finish( runner_name=runner_name, status="failure", diff --git a/uni_agent/gateway/manager.py b/uni_agent/gateway/manager.py index b96f72278..5482da188 100644 --- a/uni_agent/gateway/manager.py +++ b/uni_agent/gateway/manager.py @@ -8,12 +8,15 @@ from __future__ import annotations import asyncio +import logging import ray from uni_agent.gateway.config import GatewayActorConfig from verl.workers.rollout.llm_server import LLMServerClient +logger = logging.getLogger(__name__) + class GatewayManager: """Owns gateway actors and routes sessions to them. @@ -73,7 +76,9 @@ def _get_gateway(self, session_id: str): return self.gateways[gateway_index], gateway_index async def create_session(self, session_id: str, **kwargs): - """Create a session on the least-loaded actor, record the route, and return its handle.""" + """Create a session; cancellation settles the remote call and aborts only its result.""" + if session_id in self._session_to_gateway_index: + raise RuntimeError(f"Session {session_id} already exists") gateway_index = self._select_gateway_index() gateway = self.gateways[gateway_index] # Reserve the slot synchronously, before the await. Sessions are created @@ -85,26 +90,53 @@ async def create_session(self, session_id: str, **kwargs): self._session_to_gateway_index[session_id] = gateway_index self.active_sessions_per_gateway[gateway_index] += 1 try: - return await gateway.create_session.remote(session_id=session_id, **kwargs) + create = asyncio.ensure_future(gateway.create_session.remote(session_id=session_id, **kwargs)) + return await asyncio.shield(create) + except asyncio.CancelledError: + cleanup = asyncio.ensure_future(self._abort_cancelled_create(session_id, gateway_index, create)) + # Ray create keeps running after caller cancellation. Every wait must + # stay shielded so repeated cancellation cannot orphan its session. + while not cleanup.done(): + try: + await asyncio.shield(cleanup) + except asyncio.CancelledError: + continue + cleanup.result() + raise except BaseException: - self.active_sessions_per_gateway[gateway_index] -= 1 - self._session_to_gateway_index.pop(session_id, None) + self._drop_route(session_id, gateway_index) raise + async def _abort_cancelled_create(self, session_id: str, gateway_index: int, create: asyncio.Future) -> None: + try: + await create + except (Exception, asyncio.CancelledError): + # A failed create owns no actor session; never abort a pre-existing one. + self._drop_route(session_id, gateway_index) + return + try: + await self.abort_session(session_id) + except (Exception, asyncio.CancelledError): + # A failed abort leaves this owner available for an explicit retry. + logger.exception("session %s: abort failed after cancelled create", session_id) + async def finalize_session(self, session_id: str): """Finalize a session on its owning actor, release the route, and return its trajectories.""" gateway, gateway_index = self._get_gateway(session_id) trajectories = await gateway.finalize_session.remote(session_id=session_id) - self._session_to_gateway_index.pop(session_id, None) - self.active_sessions_per_gateway[gateway_index] -= 1 + self._drop_route(session_id, gateway_index) return trajectories async def abort_session(self, session_id: str) -> None: """Abort a routed session on its owning actor and release the route.""" gateway, gateway_index = self._get_gateway(session_id) await gateway.abort_session.remote(session_id=session_id) - self._session_to_gateway_index.pop(session_id, None) - self.active_sessions_per_gateway[gateway_index] -= 1 + self._drop_route(session_id, gateway_index) + + def _drop_route(self, session_id: str, gateway_index: int) -> None: + # Concurrent closes must release the load slot at most once. + if self._session_to_gateway_index.pop(session_id, None) is not None: + self.active_sessions_per_gateway[gateway_index] -= 1 async def shutdown(self) -> None: """Stop owned gateway actors and clear routing state."""