# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Engine-side owner of duplex sessions: admission, lease reaping, command dispatch. ``DuplexOrchestrator`` hosts exactly one manager. The whole session (``DuplexEngineSession``) lives here, owned by one ``DuplexSessionRunner`` per session on the orchestrator loop; nothing outside the engine holds session state. Control operations (open / close / resume / touch) arrive as correlated RPC messages and answer through ``result_sink``; session commands arrive one-way as ``DuplexSessionCommandMessage`` and are pushed onto the runner's ordered mailbox after backpressure admission; everything a session emits leaves through ``output_sink`` as ``DuplexSessionEventMessage``. """ from __future__ import annotations import asyncio import concurrent.futures import time from collections.abc import Awaitable, Callable, Iterable from dataclasses import dataclass, replace from typing import TYPE_CHECKING from vllm.logger import init_logger from vllm_omni.engine.duplex.commands import AppendAudio, Commit, DuplexCommand from vllm_omni.engine.duplex.contracts import ( DuplexFence, DuplexStagePort, DuplexStageRequestContext, duplex_resource_request_belongs_to_session, duplex_resource_request_id, ) from vllm_omni.engine.duplex.events import DuplexEvent, ErrorEvent, SessionClosed, SessionExpired, error_event from vllm_omni.engine.duplex.messages import ( CloseDuplexSessionMessage, DuplexControlResultMessage, DuplexSessionCommandMessage, DuplexSessionError, DuplexSessionEventMessage, OpenDuplexSessionMessage, ResumeDuplexSessionMessage, TouchDuplexSessionMessage, ) from vllm_omni.engine.duplex.plugin import DuplexModelPlugin, DuplexRuntimeConfigError, validate_duplex_plugin_sampling from vllm_omni.engine.duplex.session.engine_session import DuplexEngineSession, DuplexFenceMismatchError from vllm_omni.engine.duplex.session.lease import DuplexLeaseActivity, DuplexLeaseConfig, DuplexLeaseState from vllm_omni.engine.duplex.turn_detection import SileroVADBackendProvider if TYPE_CHECKING: import janus from vllm.config import ModelConfig from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig from vllm_omni.engine.duplex.session.runner import DuplexSessionRunner from vllm_omni.engine.messages import EngineQueueMessage logger = init_logger(__name__) @dataclass(frozen=True) class _PendingRequestCleanup: session_id: str lease_generation: int request_ids: tuple[str, ...] abort: bool @dataclass(frozen=True) class _PendingSessionCleanup: """A close/expiry whose stage cleanup has not succeeded yet (retried by the reaper).""" kind: str session_id: str lease_generation: int reason: str submitted_request_ids: tuple[str, ...] reserved_request_ids: tuple[str, ...] class DuplexSessionManager: _CONTROL_TYPES = ( OpenDuplexSessionMessage, CloseDuplexSessionMessage, ResumeDuplexSessionMessage, TouchDuplexSessionMessage, ) _MESSAGE_TYPES = (*_CONTROL_TYPES, DuplexSessionCommandMessage) def __init__( self, *, plugin: DuplexModelPlugin, stage_port: DuplexStagePort, output_sink: janus.AsyncQueue[EngineQueueMessage], result_sink: janus.AsyncQueue[EngineQueueMessage], runtime_config: DuplexSessionRuntimeConfig, model_config: ModelConfig | None, clock: Callable[[], float] | None = None, executor: concurrent.futures.ThreadPoolExecutor | None = None, ) -> None: validate_duplex_plugin_sampling(plugin, sampling_defaults=tuple(stage_port.sampling_defaults())) self.plugin = plugin self.stage_port = stage_port self.model_config = model_config self.runtime_config = runtime_config self._output_sink = output_sink self._result_sink = result_sink self._clock = clock or time.monotonic self.executor = executor or concurrent.futures.ThreadPoolExecutor( max_workers=4, thread_name_prefix="duplex-session", ) self._owns_executor = executor is None self._lease_config = DuplexLeaseConfig( idle_ttl_s=runtime_config.idle_ttl_s, disconnect_grace_s=runtime_config.disconnect_grace_s, ) #: One Silero backend for the whole engine: the ONNX graph is stateless #: per call, so every session's detector can share it. Resolved lazily, #: on the first session that actually turns server VAD on. self.vad_backend_provider = SileroVADBackendProvider( model_path=getattr(runtime_config, "server_vad_model_path", None) ) self.runners: dict[str, DuplexSessionRunner] = {} #: Sessions whose close/expiry began but whose stage cleanup has not finalized; #: they keep their admission slot until finalized. self._closing: dict[str, _PendingSessionCleanup] = {} #: Session ids whose ``open`` passed the capacity check and is still awaiting the plugin. self._admitting: set[str] = set() #: Opens whose caller gave up (control RPC timeout / cancellation) before #: admission finished. The open completes anyway -- it is already past #: the capacity check -- so it must tear itself down on arrival instead #: of holding a slot and a stage resource until idle expiry. self._abandoned_opens: set[str] = set() self._request_index: dict[str, str] = {} self._session_control_tails: dict[str, asyncio.Task[None]] = {} self._dispatched_control_tasks: set[asyncio.Task[None]] = set() self._pending_request_cleanups: dict[tuple[str, int], _PendingRequestCleanup] = {} self._request_cleanup_tasks: dict[tuple[str, int], asyncio.Task[None]] = {} self._request_cleanups_in_progress: set[tuple[str, int]] = set() #: Sessions retained for cleanup retry after their runner is gone. self._session_snapshots: dict[str, DuplexEngineSession] = {} # ------------------------------------------------------------------ # # Views # # ------------------------------------------------------------------ # def active_count(self) -> int: return len(self.runners) def get(self, session_id: str) -> DuplexEngineSession | None: runner = self.runners.get(session_id) return runner.session if runner is not None else None def sessions(self) -> dict[str, DuplexEngineSession]: return {session_id: runner.session for session_id, runner in self.runners.items()} def runner_for_request_id(self, request_id: str) -> DuplexSessionRunner | None: session_id = self._request_index.get(request_id) if session_id is None: for candidate_id in self.runners: if duplex_resource_request_belongs_to_session(request_id, candidate_id): session_id = candidate_id break if session_id is None: return None return self.runners.get(session_id) def register_request(self, request_id: str, session_id: str) -> None: self._request_index[request_id] = session_id def unregister_request(self, request_id: str) -> None: self._request_index.pop(request_id, None) def _unregister_session_requests(self, session_id: str) -> None: for request_id in [rid for rid, sid in self._request_index.items() if sid == session_id]: self._request_index.pop(request_id, None) # ------------------------------------------------------------------ # # Message routing # # ------------------------------------------------------------------ # def accepts(self, message: object) -> bool: return isinstance(message, self._MESSAGE_TYPES) def dispatch(self, message: object) -> None: """Route one engine request-queue message without blocking the request handler.""" if isinstance(message, DuplexSessionCommandMessage): self._dispatch_command(message) return if not isinstance(message, self._CONTROL_TYPES): raise TypeError(f"Unsupported duplex control message: {type(message).__name__}") session_id = message.session_id async def handle_message() -> None: await self.handle(message) self._run_control(session_id, f"duplex-control-{session_id}-{message.control_id}", handle_message) def _run_control( self, session_id: str, name: str, operation: Callable[[], Awaitable[None]], ) -> asyncio.Task[None]: """Run one control operation after the session's previous one (per-session order).""" predecessor = self._session_control_tails.get(session_id) async def run_ordered() -> None: if predecessor is not None: try: await predecessor except BaseException: pass await operation() task = asyncio.create_task(run_ordered(), name=name) self._session_control_tails[session_id] = task self._track_task(task) def discard(completed: asyncio.Task[None]) -> None: if self._session_control_tails.get(session_id) is completed: self._session_control_tails.pop(session_id, None) task.add_done_callback(discard) return task def _track_task(self, task: asyncio.Task[None]) -> None: """Keep a background task alive until it finishes and log an escaped exception.""" self._dispatched_control_tasks.add(task) def done(completed: asyncio.Task[None]) -> None: self._dispatched_control_tasks.discard(completed) if not completed.cancelled() and completed.exception() is not None: logger.error("duplex control task %s failed: %r", completed.get_name(), completed.exception()) task.add_done_callback(done) async def handle(self, message: object) -> None: if isinstance(message, OpenDuplexSessionMessage): await self.open(message) elif isinstance(message, CloseDuplexSessionMessage): await self.close(message) elif isinstance(message, ResumeDuplexSessionMessage): await self.resume(message) elif isinstance(message, TouchDuplexSessionMessage): await self.touch(message) elif isinstance(message, DuplexSessionCommandMessage): self._dispatch_command(message) else: raise TypeError(f"Unsupported duplex control message: {type(message).__name__}") def _dispatch_command(self, message: DuplexSessionCommandMessage) -> None: """Admit one command: identity check, backpressure, then the runner's ordered mailbox. Backpressure contract shared with the runner: a ``Commit`` that passes admission holds one pending-turn reservation and an ``AppendAudio`` holds its wire-size byte reservation until the runner dequeues the command (the runner then re-reserves the decoded size around its PCM reservation), so the mailbox never holds more bytes than the session limit even while the worker is busy. """ runner = self.runners.get(message.session_id) command = message.command if runner is None: self._emit_raw( message.session_id, self._error_event( "unknown_session", f"Unknown or closed duplex session: {message.session_id}", command=command, ), ) return session = runner.session if runner.closing: self.emit( session, self._error_event( "session_closed", f"Duplex session is closing: {session.session_id}", command=command ), ) return if isinstance(command, AppendAudio): limit = int(self.runtime_config.max_pending_input_bytes_per_session) if not session.reserve_input_bytes(len(command.audio), limit=limit): self.emit( session, self._error_event( "input_backpressure", "Duplex session has too many pending input bytes", command=command, ), ) return elif isinstance(command, Commit): if not session.reserve_pending_turn(limit=int(self.runtime_config.max_pending_turns_per_session)): self.emit( session, self._error_event( "input_backpressure", "Duplex session has too many pending input turns", command=command, ), ) return runner.submit(command) @staticmethod def _error_event(code: str, message: str, *, command: DuplexCommand | None = None) -> ErrorEvent: return error_event(code, message, event_id=command.event_id if command is not None else None) # ------------------------------------------------------------------ # # Emission # # ------------------------------------------------------------------ # def emit(self, session: DuplexEngineSession, event: DuplexEvent) -> None: """Bind the session identity to one typed event and push it to the engine output queue.""" self._emit_raw(session.session_id, event, epoch=session.epoch) def _emit_raw( self, session_id: str, event: DuplexEvent, *, epoch: int | None = None, ) -> None: event = replace(event, session_id=session_id, epoch=epoch) message = DuplexSessionEventMessage(session_id=session_id, event=event) put_nowait = getattr(self._output_sink, "put_nowait", None) if callable(put_nowait): put_nowait(message) else: # pragma: no cover - sink without put_nowait asyncio.ensure_future(self._output_sink.put(message)) async def _put_result( self, message: ( OpenDuplexSessionMessage | CloseDuplexSessionMessage | ResumeDuplexSessionMessage | TouchDuplexSessionMessage ), *, operation: str, ok: bool, session: DuplexEngineSession | None = None, error: BaseException | None = None, ) -> None: error_code, error_message, error_retryable = ( self._control_error(error) if error is not None else (None, None, False) ) result = DuplexControlResultMessage( control_id=message.control_id, operation=operation, session_id=message.session_id, ok=ok, lease_generation=session.lease_generation if session is not None else None, capabilities=session.capabilities if session is not None and ok else None, public_session=session.as_public_dict() if session is not None and ok else None, error_code=error_code, error_message=error_message, error_retryable=error_retryable, ) await self._result_sink.put(result) @staticmethod def _control_error(error: BaseException) -> tuple[str, str, bool]: """Map an exception to ``(code, message, retryable)`` for a control result.""" message = str(error) retryable = False if isinstance(error, DuplexSessionError): code = error.code retryable = error.retryable elif isinstance(error, DuplexRuntimeConfigError): code = error.code elif isinstance(error, DuplexFenceMismatchError): code = "stale_fence" elif isinstance(error, KeyError): code = "not_found" elif isinstance(error, TypeError | ValueError): code = "invalid_argument" elif isinstance(error, TimeoutError): code = "timeout" retryable = True else: code = "failed_precondition" return code, message, retryable # ------------------------------------------------------------------ # # Sampling / stage request helpers shared with the runner # # ------------------------------------------------------------------ # def sampling_params_for(self, session: DuplexEngineSession) -> tuple[object, ...]: defaults = tuple(self.stage_port.sampling_defaults()) configured = self.plugin.configure_sampling_params( runtime_config=dict(session.runtime_config), defaults=defaults, ) if not isinstance(configured, tuple): raise TypeError("duplex plugin must return sampling parameters as a tuple") if len(configured) != len(defaults): raise ValueError("duplex plugin must return one sampling parameter per stage") return configured @staticmethod def stage_request_id(fence: DuplexFence, *, stage_id: int, resumable: bool = True) -> str: role = f"stage{stage_id}" if not resumable: role += f"-turn{fence.turn_id}" return duplex_resource_request_id(fence, role) def ensure_stage_request( self, session: DuplexEngineSession, *, stage_id: int, fence: DuplexFence | None = None, ) -> DuplexStageRequestContext | None: """Reserve the session's stage request id and register it with the stage port.""" if stage_id >= self.stage_port.stage_count: return None effective_fence = fence or session.fence request_id = self.stage_request_id( effective_fence, stage_id=stage_id, resumable=session.capabilities.supports_core_resumable_request ) session.reserve_stage_request(stage_id, request_id, fence=effective_fence) context = DuplexStageRequestContext( request_id=request_id, session_id=session.session_id, fence=effective_fence, stage_id=stage_id, final_stage_id=self.stage_port.stage_count - 1, config_generation=session.config_generation, sampling_params=self.sampling_params_for(session), session_config=session.config.as_dict(), runtime_config=session.runtime_config, ) self.stage_port.ensure_request(context) self.register_request(request_id, session.session_id) return context # ------------------------------------------------------------------ # # Control operations # # ------------------------------------------------------------------ # def _admission_count(self) -> int: return len(set(self.runners) | set(self._closing) | self._admitting) async def open(self, message: OpenDuplexSessionMessage) -> None: from vllm_omni.engine.duplex.session.runner import DuplexSessionRunner session_id = message.session_id session: DuplexEngineSession | None = None runner: DuplexSessionRunner | None = None # Only the open that put the id in ``_admitting`` may take it out again: # a duplicate open must not release the slot the first one is holding. holds_admission_slot = False try: if session_id in self.runners or session_id in self._closing or session_id in self._admitting: raise DuplexSessionError(f"Duplex session already exists: {session_id}", code="session_exists") max_sessions = int(self.runtime_config.max_sessions) if self._admission_count() >= max_sessions: raise DuplexSessionError( f"duplex_session_capacity_exhausted: limit={max_sessions}", code="resource_exhausted", retryable=True, ) # Hold the slot across the plugin await: opens of different # sessions run concurrently and must not all pass the check above. self._admitting.add(session_id) holds_admission_slot = True config = message.session_config self.plugin.validate_client_extra_body(config.extra_body) try: runtime_config = await self.plugin.prepare_runtime_config(config, model_config=self.model_config) except DuplexRuntimeConfigError: raise except ValueError as exc: raise DuplexRuntimeConfigError(str(exc)) from exc capabilities = self.plugin.capabilities(max_sessions=max_sessions) session = DuplexEngineSession( session_id=session_id, config=config, capabilities=capabilities, lease=DuplexLeaseState(config=self._lease_config, generation=0, last_activity=self._clock()), _clock=self._clock, _runtime_config=dict(runtime_config), ) # Validates the plugin's sampling policy for this runtime config before admission. self.sampling_params_for(session) session.model_state = self.plugin.create_session_state() runner = DuplexSessionRunner( session=session, plugin=self.plugin, stage_port=self.stage_port, manager=self, model_config=self.model_config, ) self.runners[session_id] = runner # Reserve the Stage0 request resource atomically with admission. self.ensure_stage_request(session, stage_id=0) runner.start() await self._put_result(message, operation="open", ok=True, session=session) if session_id in self._abandoned_opens: # The caller timed out while we were admitting; hand the slot # and the stage resource straight back. self._abandoned_opens.discard(session_id) await self._close_runner(runner, kind="close", reason="open_abandoned") except Exception as exc: error_code, _, _ = self._control_error(exc) if error_code in {"resource_exhausted", "session_exists"}: logger.info("open_duplex_session rejected: %s", exc) else: logger.exception("open_duplex_session failed: %s", exc) if runner is not None and self.runners.get(session_id) is runner: self.runners.pop(session_id, None) try: await runner.shutdown() except Exception: logger.exception("duplex open rollback: runner shutdown failed for %s", session_id) if session is not None: reserved = session.release_all_requests() if reserved: try: await self.stage_port.cleanup(list(reserved)) except Exception: logger.warning("duplex open rollback: request cleanup pending for %s", session_id) self._unregister_session_requests(session_id) await self._put_result(message, operation="open", ok=False, error=exc) finally: if holds_admission_slot: self._admitting.discard(session_id) self._abandoned_opens.discard(session_id) def _require_runner(self, session_id: str) -> DuplexSessionRunner: runner = self.runners.get(session_id) if runner is None: raise DuplexSessionError(f"Unknown duplex session: {session_id}", code="unknown_session") if runner.closing: raise DuplexSessionError(f"Duplex session is closing: {session_id}", code="session_closed") return runner async def close(self, message: CloseDuplexSessionMessage) -> None: session_id = message.session_id runner = self.runners.get(session_id) if runner is None: if session_id in self._admitting: # The open is still awaiting the plugin. It cannot be cancelled # from here, so leave a note it checks before returning. self._abandoned_opens.add(session_id) # Already closing, or closed and forgotten: ids are never reused # (D16), so a repeated close is idempotent rather than an error. await self._put_result(message, operation="close", ok=True) return session = runner.session try: await self._close_runner(runner, kind="close", reason=message.reason) except Exception as exc: # The session is closed either way; only the stage cleanup is # outstanding and the reaper retries it while the slot stays held. logger.warning("close_duplex_session: stage cleanup for %s remains pending: %s", session_id, exc) await self._put_result(message, operation="close", ok=True, session=session) def close_from_runner(self, runner: DuplexSessionRunner, reason: str) -> None: """Finish a close the runner started from its command stream (``session.close``). Runs on the orchestrator loop as a tracked task (the runner's worker cannot await its own teardown): stage requests are aborted, the admission slot is released and ``session.closed`` is emitted last. """ session_id = runner.session.session_id if self.runners.get(session_id) is not runner: return async def close_runner() -> None: if self.runners.get(session_id) is runner: await self._close_runner(runner, kind="close", reason=reason) # Chained on the session's control tail so a control RPC issued after # the wire close observes the closed session, never a half-torn one. self._run_control(session_id, f"duplex-close-{session_id}", close_runner) async def _close_runner(self, runner: DuplexSessionRunner, *, kind: str, reason: str) -> None: """Begin an irreversible close, tear down the runner, then release stage resources. The session keeps its admission slot (``_closing``) until stage cleanup succeeded; a failed cleanup is retried by the reaper. For an explicit close the ``session.closed`` event is emitted only after the cleanup attempt, so a client that sees it can open a replacement session at once (the runner used to emit it before the stage requests were aborted, which let a prompt reopen hit ``resource_exhausted``). """ session = runner.session submitted = tuple(session.resource_request_ids(submitted=True)) reserved = tuple(session.resource_request_ids(submitted=False)) session.begin_close(reason=reason) pending = _PendingSessionCleanup( kind=kind, session_id=session.session_id, lease_generation=session.lease_generation, reason=reason, submitted_request_ids=submitted, reserved_request_ids=reserved, ) self._closing[session.session_id] = pending if self.runners.get(session.session_id) is runner: self.runners.pop(session.session_id, None) emit_after_cleanup = kind in {"close", "expired"} and not runner.closed_emitted try: if kind == "expired": await runner.expire(reason, emit_expired=not emit_after_cleanup) else: await runner.close(reason, emit_closed=not emit_after_cleanup) except Exception: logger.exception("duplex runner %s failed for session %s", kind, session.session_id) session.close() try: await self._finalize_pending_cleanup(pending, session) finally: # ``emit_after_cleanup`` was decided before the teardown awaits; a # runtime close racing them may have emitted the terminal already. if emit_after_cleanup and not runner.closed_emitted: runner.mark_closed_emitted() terminal: DuplexEvent = ( SessionExpired(reason=reason) if kind == "expired" else SessionClosed(reason=reason) ) self.emit(session, terminal) async def _finalize_pending_cleanup( self, pending: _PendingSessionCleanup, session: DuplexEngineSession | None, ) -> None: if pending.submitted_request_ids: await self.stage_port.cleanup(list(pending.submitted_request_ids), abort=True) if pending.reserved_request_ids: await self.stage_port.cleanup(list(pending.reserved_request_ids)) if session is not None: session.release_all_requests() self._unregister_session_requests(pending.session_id) if self._closing.get(pending.session_id) is pending: self._closing.pop(pending.session_id, None) self._session_snapshots.pop(pending.session_id, None) async def resume(self, message: ResumeDuplexSessionMessage) -> None: session: DuplexEngineSession | None = None try: runner = self._require_runner(message.session_id) session = runner.session try: session.resume_lease(expected_lease_generation=message.expected_lease_generation) except ValueError as exc: raise DuplexSessionError(str(exc), code="session_resume_conflict") from exc await self._put_result(message, operation="resume", ok=True, session=session) except Exception as exc: logger.exception("resume_duplex_session failed: %s", exc) await self._put_result(message, operation="resume", ok=False, session=session, error=exc) async def touch(self, message: TouchDuplexSessionMessage) -> None: session: DuplexEngineSession | None = None try: runner = self._require_runner(message.session_id) session = runner.session activity = DuplexLeaseActivity(message.activity) if activity is DuplexLeaseActivity.DETACH: session.detach_lease() else: session.touch_lease(activity) await self._put_result(message, operation="touch", ok=True, session=session) except DuplexSessionError as exc: # A lease touch racing the session's own close is routine: the # caller learns the session is gone from the result, not a fault. logger.debug("touch_duplex_session skipped: %s", exc) await self._put_result(message, operation="touch", ok=False, session=session, error=exc) except Exception as exc: logger.exception("touch_duplex_session failed: %s", exc) await self._put_result(message, operation="touch", ok=False, session=session, error=exc) # ------------------------------------------------------------------ # # Lease expiry / reaper # # ------------------------------------------------------------------ # async def reaper_loop(self, shutdown_event: asyncio.Event) -> None: interval = float(self.runtime_config.reaper_interval_s) while not shutdown_event.is_set(): try: await asyncio.wait_for(shutdown_event.wait(), timeout=interval) except TimeoutError: try: await self.reap_expired() except Exception: logger.exception("[DuplexSessionManager] expiry cleanup failed; retrying on next tick") async def reap_expired(self, now: float | None = None) -> int: completed = 0 effective_now = self._clock() if now is None else now # Retry request cleanups recorded by the orchestrator error paths. for key, request_cleanup in list(self._pending_request_cleanups.items()): if key in self._request_cleanups_in_progress: continue try: await self._complete_request_cleanup(key, request_cleanup) except Exception as exc: logger.warning( "duplex request cleanup remains pending for session %s: %s", request_cleanup.session_id, exc, ) continue completed += 1 # Retry closes/expiries whose stage cleanup failed. for session_id, session_cleanup in list(self._closing.items()): if session_id in self.runners: continue if session_cleanup.kind == "request_cleanup" and any( key[0] == session_id for key in self._pending_request_cleanups ): # The stage cleanup is owned by the request-cleanup path (orchestrator # or the retry above); ``finalize_closed_sessions`` releases the slot. continue try: await self._finalize_pending_cleanup(session_cleanup, self._session_snapshots.get(session_id)) except Exception as exc: logger.warning( "duplex %s cleanup remains pending for session %s: %s", session_cleanup.kind, session_id, exc, ) continue completed += 1 # Expire leases. for session_id, runner in list(self.runners.items()): lease = runner.session.lease if lease.disconnect_grace_expired(effective_now): reason = "disconnect_grace_expired" elif lease.idle_expired(effective_now): reason = "idle_ttl_expired" else: continue self._session_snapshots[session_id] = runner.session try: await self._close_runner(runner, kind="expired", reason=reason) except Exception as exc: logger.warning("duplex expiry cleanup remains pending for session %s: %s", session_id, exc) continue completed += 1 return completed # ------------------------------------------------------------------ # # Request-triggered cleanup (orchestrator error paths) # # ------------------------------------------------------------------ # def close_sessions_for_request_ids( self, request_ids: list[str], *, abort: bool = False, cleanup_in_progress: bool = False, ) -> dict[str, list[str]]: request_id_set = set(request_ids) closed: dict[str, list[str]] = {} for session_id, runner in list(self.runners.items()): session = runner.session stale = session.resource_request_ids() if request_id_set.isdisjoint(stale): continue if not session.begin_close(reason="request_cleanup"): continue closed[session_id] = stale key = (session_id, session.lease_generation) existing = self._pending_request_cleanups.get(key) merged = tuple(dict.fromkeys([*(existing.request_ids if existing is not None else ()), *stale])) self._pending_request_cleanups[key] = _PendingRequestCleanup( session_id=session_id, lease_generation=session.lease_generation, request_ids=merged, abort=abort or (existing.abort if existing is not None else False), ) if cleanup_in_progress: self._request_cleanups_in_progress.add(key) # Tear the runner down in the background; the orchestrator owns the stage cleanup. self._session_snapshots[session_id] = session self._closing[session_id] = _PendingSessionCleanup( kind="request_cleanup", session_id=session_id, lease_generation=session.lease_generation, reason="request_cleanup", submitted_request_ids=(), reserved_request_ids=(), ) self.runners.pop(session_id, None) task = asyncio.create_task(runner.expire("request_cleanup"), name=f"duplex-request-cleanup-{session_id}") self._track_task(task) return closed def defer_request_cleanups(self, session_ids: Iterable[str]) -> None: session_id_set = set(session_ids) active_keys = {key for key in self._request_cleanups_in_progress if key[0] in session_id_set} self._request_cleanups_in_progress.difference_update(active_keys) def finalize_closed_sessions(self, session_ids: Iterable[str]) -> None: session_id_set = set(session_ids) for key in list(self._pending_request_cleanups): if key[0] not in session_id_set: continue self._pending_request_cleanups.pop(key, None) self._request_cleanups_in_progress.discard(key) for session_id in session_id_set: session = self._session_snapshots.pop(session_id, None) if session is not None and session.lease.terminal_reason is not None: session.release_all_requests() session.close() self._closing.pop(session_id, None) self._unregister_session_requests(session_id) async def _complete_request_cleanup( self, key: tuple[str, int], pending: _PendingRequestCleanup, ) -> None: task = self._request_cleanup_tasks.get(key) if task is not None and task.done(): self._request_cleanup_tasks.pop(key, None) task = None if task is None: task = asyncio.create_task( self._run_request_cleanup(key, pending), name=f"duplex-request-cleanup-{pending.session_id}", ) self._request_cleanup_tasks[key] = task def discard(completed: asyncio.Task[None]) -> None: if self._request_cleanup_tasks.get(key) is completed: self._request_cleanup_tasks.pop(key, None) task.add_done_callback(discard) await asyncio.shield(task) async def _run_request_cleanup(self, key: tuple[str, int], pending: _PendingRequestCleanup) -> None: await self.stage_port.cleanup(list(pending.request_ids), abort=pending.abort) if self._pending_request_cleanups.get(key) is pending: self._pending_request_cleanups.pop(key, None) self._request_cleanups_in_progress.discard(key) self.finalize_closed_sessions([pending.session_id]) # ------------------------------------------------------------------ # # Shutdown # # ------------------------------------------------------------------ # async def shutdown(self) -> None: runners = list(self.runners.values()) self.runners.clear() for runner in runners: try: await runner.shutdown() except Exception: logger.exception("duplex runner shutdown failed for %s", runner.session.session_id) tasks = tuple(self._dispatched_control_tasks | set(self._request_cleanup_tasks.values())) for task in tasks: task.cancel() if tasks: await asyncio.gather(*tasks, return_exceptions=True) self._session_control_tails.clear() self._closing.clear() self._session_snapshots.clear() self._request_index.clear() if self._owns_executor: self.executor.shutdown(wait=False, cancel_futures=True) __all__ = ["DuplexSessionManager"]