Skip to content

Clean up WebRTCSession resources after startup failures #2808

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
voropaevv wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Clean up WebRTCSession resources after startup failures #2808

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
213 changes: 161 additions & 52 deletions inference_sdk/webrtc/session.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -222,6 +222,7 @@ def __init__(
# Internal state
self._loop: Optional[asyncio.AbstractEventLoop] = None
self._loop_thread: Optional[threading.Thread] = None
self._startup_task: Optional["asyncio.Task[None]"] = None
self._pc: Optional["RTCPeerConnection"] = None
# In model mode queue items carry an extra raw-predictions-dict element:
# (frame, data, metadata). In default mode: (frame, metadata).
Expand Down Expand Up @@ -265,53 +266,117 @@ def __init__(
def _init_connection(self) -> None:
"""Initialize event loop, thread, and WebRTC connection."""
# Start event loop in background thread
self._loop = asyncio.new_event_loop()
loop = asyncio.new_event_loop()
self._loop = loop

def _run(loop: asyncio.AbstractEventLoop) -> None:
asyncio.set_event_loop(loop)
loop.run_forever()
try:
loop.run_forever()
finally:
try:
try:
pending_tasks = asyncio.all_tasks(loop)
for task in pending_tasks:
task.cancel()
if pending_tasks:
loop.run_until_complete(
asyncio.gather(*pending_tasks, return_exceptions=True)
)
except Exception:
logger.exception("Failed to drain WebRTC event-loop tasks")

self._loop_thread = threading.Thread(
target=_run, args=(self._loop,), daemon=True
)
self._loop_thread.start()
try:
loop.run_until_complete(loop.shutdown_asyncgens())
except Exception:
logger.exception("Failed to shut down WebRTC async generators")
finally:
try:
loop.close()
finally:
asyncio.set_event_loop(None)

loop_thread = threading.Thread(target=_run, args=(loop,), daemon=True)
self._loop_thread = loop_thread
try:
loop_thread.start()
except BaseException:
if not loop_thread.is_alive():
loop.close()
self._loop = None
self._loop_thread = None
raise

# Initialize WebRTC connection
fut = asyncio.run_coroutine_threadsafe(self._init(), self._loop)
init_coro = self._run_startup()
try:
fut = asyncio.run_coroutine_threadsafe(init_coro, loop)
except BaseException:
init_coro.close()
raise

try:
fut.result()
except requests.exceptions.HTTPError as e:
if e.response.status_code == 404:
raise RuntimeError(
f"WebRTC endpoint not found at {self._api_url}/initialise_webrtc_worker.\n"
f"This API URL may not support WebRTC streaming.\n"
f"Troubleshooting:\n"
f" - For self-hosted inference, ensure the server is started with WebRTC enabled\n"
f" - For Roboflow Cloud, use a dedicated inference server URL (not serverless.roboflow.com)\n"
f" - Verify the --api-url parameter points to the correct server\n"
f"Response: {e.response.text}"
) from e
else:
except BaseException as error:
if not fut.done():
fut.cancel()

if isinstance(error, requests.exceptions.HTTPError):
if error.response.status_code == 404:
raise RuntimeError(
f"WebRTC endpoint not found at {self._api_url}/initialise_webrtc_worker.\n"
f"This API URL may not support WebRTC streaming.\n"
f"Troubleshooting:\n"
f" - For self-hosted inference, ensure the server is started with WebRTC enabled\n"
f" - For Roboflow Cloud, use a dedicated inference server URL (not serverless.roboflow.com)\n"
f" - Verify the --api-url parameter points to the correct server\n"
f"Response: {error.response.text}"
) from error
raise RuntimeError(
f"Failed to initialize WebRTC session (HTTP {e.response.status_code}).\n"
f"Failed to initialize WebRTC session (HTTP {error.response.status_code}).\n"
f"API URL: {self._api_url}\n"
f"Error: {e}\n"
f"Response: {e.response.text}"
) from e
except Exception as e:
f"Error: {error}\n"
f"Response: {error.response.text}"
) from error

if not isinstance(error, Exception):
raise

raise RuntimeError(
f"Failed to initialize WebRTC session: {e.__class__.__name__}: {e}\n"
f"Failed to initialize WebRTC session: "
f"{error.__class__.__name__}: {error}\n"
f"API URL: {self._api_url}"
) from e
) from error

async def _run_startup(self) -> None:
"""Track the loop-owned startup task until teardown observes it."""
self._startup_task = asyncio.current_task()
await self._init()

def _ensure_started(self) -> None:
"""Ensure connection is started (thread-safe, idempotent)."""
with self._state_lock:
if self._state == SessionState.NOT_STARTED:
self._state = SessionState.STARTED
self._init_connection()
elif self._state == SessionState.CLOSED:
raise RuntimeError("Cannot use closed WebRTCSession")
startup_failed = False
try:
with self._state_lock:
if self._state == SessionState.NOT_STARTED:
self._state = SessionState.STARTED
try:
self._init_connection()
except BaseException:
self._state = SessionState.CLOSED
startup_failed = True
raise
elif self._state == SessionState.CLOSED:
raise RuntimeError("Cannot use closed WebRTCSession")
except BaseException:
if startup_failed:
try:
self.close()
except (Exception, asyncio.CancelledError):
logger.exception(
"Failed to clean up WebRTC session after startup failure"
)
raise

def _parse_video_metadata(
self,
Expand Down Expand Up @@ -401,40 +466,85 @@ def _close_in_background(self) -> None:
except Exception:
logger.exception("Failed to close WebRTC session")

async def _cleanup_async_resources(self) -> None:
"""Wait for startup to stop before reading and closing its resources."""
startup_task = self._startup_task
try:
if startup_task is not None and startup_task is not asyncio.current_task():
if not startup_task.done():
startup_task.cancel()
try:
await startup_task
except asyncio.CancelledError:
pass
except Exception:
# The startup caller owns reporting its original exception.
pass
finally:
self._startup_task = None
try:
if self._pc is not None:
await self._pc.close()
finally:
if self._source is not None:
await self._source.cleanup()

def _close_resources(self) -> None:
"""Close session resources once and signal waiting callers."""
try:
with self._state_lock:
self._state = SessionState.CLOSED
loop = self._loop
loop_thread = self._loop_thread

# Signal video iterator to stop by putting None sentinel
try:
self._video_queue.put_nowait(None)
except Exception:
pass # Queue might be full, but that's okay

# Cleanup resources (nested finally ensures all cleanup steps execute)
# Cleanup resources before stopping their owner event loop.
try:
# Close peer connection
if self._loop and self._pc:
asyncio.run_coroutine_threadsafe(
self._pc.close(), self._loop
).result()
if (
loop is not None
and not loop.is_closed()
and loop_thread is not None
and loop_thread.is_alive()
and (
self._startup_task is not None
or self._pc is not None
or self._source is not None
)
):
cleanup_coro = self._cleanup_async_resources()
try:
cleanup_future = asyncio.run_coroutine_threadsafe(
cleanup_coro, loop
)
except BaseException:
cleanup_coro.close()
raise
cleanup_future.result()
finally:
# Stop event loop and join thread
try:
# Cleanup source (webcam, video file, etc.)
if self._loop and self._source:
asyncio.run_coroutine_threadsafe(
self._source.cleanup(), self._loop
).result()
if loop is not None and not loop.is_closed():
try:
loop.call_soon_threadsafe(loop.stop)
except RuntimeError:
if not loop.is_closed():
raise
finally:
# Stop event loop and join thread
if self._loop:
self._loop.call_soon_threadsafe(self._loop.stop)
if self._loop_thread:
self._loop_thread.join(
timeout=WEBRTC_EVENT_LOOP_SHUTDOWN_TIMEOUT
)
if (
loop_thread is not None
and loop_thread is not threading.current_thread()
):
loop_thread.join(timeout=WEBRTC_EVENT_LOOP_SHUTDOWN_TIMEOUT)
if loop_thread.is_alive():
logger.warning(
"WebRTC event loop thread did not stop "
f"within {WEBRTC_EVENT_LOOP_SHUTDOWN_TIMEOUT}s"
)
finally:
self._close_done.set()

Expand Down Expand Up @@ -969,6 +1079,7 @@ async def _init(self) -> None:
turn_config = await self._get_turn_config()

pc = RTCPeerConnection(configuration=turn_config)
self._pc = pc
relay = MediaRelay()

# Monitor ICE connection state for failures
Expand Down Expand Up @@ -1275,5 +1386,3 @@ def _on_data_message(message: Any) -> None: # noqa: ANN401
# Start video file upload if applicable
if isinstance(self._source, VideoFileSource):
asyncio.ensure_future(self._source.start_upload())

self._pc = pc
Loading