wehub-resource-sync
257 lines
9.0 KiB
Python
257 lines
9.0 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import socket
|
|
import sys
|
|
from collections.abc import Callable, Coroutine
|
|
from types import FrameType
|
|
|
|
from ..log import logger
|
|
from ..utils import aio, log_exceptions, time_ms
|
|
from .channel import Message, arecv_message, asend_message, recv_message, send_message
|
|
from .proto import (
|
|
IPC_MESSAGES,
|
|
InitializeRequest,
|
|
InitializeResponse,
|
|
PingRequest,
|
|
PongResponse,
|
|
)
|
|
|
|
|
|
class _ProcClient:
|
|
def __init__(
|
|
self,
|
|
mp_cch: socket.socket,
|
|
log_cch: socket.socket | None,
|
|
initialize_fnc: Callable[[InitializeRequest, _ProcClient], None],
|
|
main_task_fnc: Callable[[aio.ChanReceiver[Message]], Coroutine[None, None, None]],
|
|
) -> None:
|
|
self._mp_cch = mp_cch
|
|
self._initialize_fnc = initialize_fnc
|
|
self._main_task_fnc = main_task_fnc
|
|
self._initialized = False
|
|
|
|
def initialize(self) -> None:
|
|
try:
|
|
cch = aio.duplex_unix._Duplex.open(self._mp_cch)
|
|
first_req = recv_message(cch, IPC_MESSAGES)
|
|
|
|
assert isinstance(first_req, InitializeRequest), (
|
|
"first message must be proto.InitializeRequest"
|
|
)
|
|
|
|
self._init_req = first_req
|
|
try:
|
|
self._initialize_fnc(self._init_req, self)
|
|
send_message(cch, InitializeResponse())
|
|
except Exception as e:
|
|
send_message(cch, InitializeResponse(error=str(e)))
|
|
raise
|
|
|
|
self._initialized = True
|
|
cch.detach()
|
|
except aio.duplex_unix.DuplexClosed as e:
|
|
raise RuntimeError("failed to initialize proc_client") from e
|
|
|
|
def run(self) -> None:
|
|
if not self._initialized:
|
|
raise RuntimeError("proc_client not initialized")
|
|
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
loop.set_debug(self._init_req.asyncio_debug)
|
|
loop.slow_callback_duration = 0.1 # 100ms
|
|
|
|
try:
|
|
self._task = loop.create_task(self._monitor_task(), name="proc_client_main")
|
|
while not self._task.done():
|
|
try:
|
|
loop.run_until_complete(self._task)
|
|
except KeyboardInterrupt:
|
|
# ignore the keyboard interrupt, we handle the process shutdown ourselves on the worker process # noqa: E501
|
|
# (See proto.ShutdownRequest)
|
|
pass
|
|
|
|
except KeyboardInterrupt:
|
|
pass
|
|
finally:
|
|
loop.run_until_complete(loop.shutdown_default_executor())
|
|
|
|
async def send(self, msg: Message) -> None:
|
|
await asend_message(self._acch, msg)
|
|
|
|
async def _monitor_task(self) -> None:
|
|
self._acch = await aio.duplex_unix._AsyncDuplex.open(self._mp_cch)
|
|
try:
|
|
exit_flag = asyncio.Event()
|
|
ping_timeout = aio.sleep(self._init_req.ping_timeout + 10)
|
|
|
|
ipc_ch = aio.Chan[Message]()
|
|
|
|
@log_exceptions(logger=logger)
|
|
async def _read_ipc_task() -> None:
|
|
while True:
|
|
try:
|
|
msg = await arecv_message(self._acch, IPC_MESSAGES)
|
|
except aio.duplex_unix.DuplexClosed:
|
|
break
|
|
|
|
with contextlib.suppress(aio.SleepFinished):
|
|
ping_timeout.reset()
|
|
|
|
if isinstance(msg, PingRequest):
|
|
await asend_message(
|
|
self._acch,
|
|
PongResponse(last_timestamp=msg.timestamp, timestamp=time_ms()),
|
|
)
|
|
|
|
ipc_ch.send_nowait(msg)
|
|
|
|
@log_exceptions(logger=logger)
|
|
async def _self_health_check() -> None:
|
|
await ping_timeout
|
|
print(
|
|
"worker process is not responding.. worker crashed?",
|
|
file=sys.stderr,
|
|
)
|
|
|
|
read_task = asyncio.create_task(_read_ipc_task(), name="ipc_read")
|
|
health_check_task: asyncio.Task[None] | None = None
|
|
if self._init_req.ping_interval > 0:
|
|
health_check_task = asyncio.create_task(_self_health_check(), name="health_check")
|
|
main_task = asyncio.create_task(
|
|
self._main_task_fnc(ipc_ch), name="main_task_entrypoint"
|
|
)
|
|
|
|
def _done_cb(_: asyncio.Task[None]) -> None:
|
|
with contextlib.suppress(asyncio.InvalidStateError):
|
|
exit_flag.set()
|
|
|
|
ipc_ch.close()
|
|
|
|
read_task.add_done_callback(_done_cb)
|
|
if health_check_task is not None:
|
|
health_check_task.add_done_callback(_done_cb)
|
|
|
|
main_task.add_done_callback(_done_cb)
|
|
|
|
await exit_flag.wait()
|
|
await aio.cancel_and_wait(read_task, main_task)
|
|
if health_check_task is not None:
|
|
await aio.cancel_and_wait(health_check_task)
|
|
|
|
finally:
|
|
await self._acch.aclose()
|
|
|
|
|
|
def _dump_stack_traces_impl() -> None:
|
|
"""Implementation of stack trace dumping (callable directly or from signal handler)."""
|
|
import asyncio
|
|
import faulthandler
|
|
import os
|
|
import tempfile
|
|
import time
|
|
import traceback
|
|
from multiprocessing import current_process
|
|
from pathlib import Path
|
|
|
|
import psutil
|
|
|
|
if os.getenv("LK_DUMP_STACK_TRACES", "0").lower() in ("0", "false", "no"):
|
|
return
|
|
|
|
dir: str = os.getenv("LK_DUMP_DIR", tempfile.gettempdir())
|
|
Path(dir).mkdir(parents=True, exist_ok=True)
|
|
|
|
with tempfile.NamedTemporaryFile(
|
|
mode="w",
|
|
dir=dir,
|
|
delete=False,
|
|
prefix=f"livekit-agents-pid-{current_process().pid}-{time.time_ns()}-",
|
|
suffix=".stacktrace",
|
|
) as f:
|
|
print(f"\n{'=' * 60}", file=f)
|
|
print(
|
|
f"Process {current_process().name} (pid {current_process().pid}) stack trace dump",
|
|
file=f,
|
|
)
|
|
print(f"{'=' * 60}\n", file=f)
|
|
|
|
faulthandler.dump_traceback(file=f, all_threads=True)
|
|
print("\n", file=f)
|
|
|
|
try:
|
|
try:
|
|
loop = asyncio.get_running_loop()
|
|
except RuntimeError:
|
|
loop = None
|
|
|
|
if loop is not None:
|
|
print("=" * 60, file=f)
|
|
print("ASYNCIO TASKS", file=f)
|
|
print("=" * 60, file=f)
|
|
|
|
tasks = asyncio.all_tasks(loop)
|
|
print(f"Total tasks: {len(tasks)}\n", file=f)
|
|
|
|
for i, task in enumerate(tasks, 1):
|
|
print(f"\n--- Task {i}/{len(tasks)} ---", file=f)
|
|
print(f"Name: {task.get_name()}", file=f)
|
|
print(f"Done: {task.done()}", file=f)
|
|
|
|
if not task.done():
|
|
print(f"Cancelled: {task.cancelled()}", file=f)
|
|
|
|
try:
|
|
stack = task.get_stack()
|
|
print(f"Stack frames: {len(stack)}", file=f)
|
|
print("Stack trace:", file=f)
|
|
for frame in stack:
|
|
traceback.print_stack(frame, limit=1, file=f)
|
|
except Exception as e:
|
|
print(f"Could not get stack: {e}", file=f)
|
|
|
|
try:
|
|
coro = task.get_coro()
|
|
print(f"Coroutine: {coro}", file=f)
|
|
if cr_frame := getattr(coro, "cr_frame", None):
|
|
print("Coroutine frame:", file=f)
|
|
traceback.print_stack(cr_frame, file=f)
|
|
except Exception as e:
|
|
print(f"Could not get coroutine: {e}", file=f)
|
|
else:
|
|
try:
|
|
exc = task.exception()
|
|
if exc:
|
|
print(f"Exception: {exc}", file=f)
|
|
print("Exception traceback:", file=f)
|
|
traceback.print_exception(type(exc), exc, exc.__traceback__, file=f)
|
|
except Exception as e:
|
|
print(f"Could not get exception: {e}", file=f)
|
|
|
|
print("", file=f)
|
|
else:
|
|
print("No asyncio event loop running", file=f)
|
|
except Exception as e:
|
|
print(f"Error dumping asyncio tasks: {e}", file=f)
|
|
traceback.print_exc(file=f)
|
|
|
|
try:
|
|
process = psutil.Process()
|
|
memory_info = process.memory_info()
|
|
memory_mb = memory_info.rss / (1024 * 1024)
|
|
|
|
print("\n" + "=" * 60, file=f)
|
|
print("MEMORY USAGE", file=f)
|
|
print("=" * 60, file=f)
|
|
print(f"RSS: {memory_mb:.2f} MB", file=f)
|
|
print(f"VMS: {memory_info.vms / (1024 * 1024):.2f} MB", file=f)
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def _dump_stack_traces(signum: int, _: FrameType | None) -> None:
|
|
"""Signal handler wrapper for _dump_stack_traces_impl."""
|
|
_dump_stack_traces_impl()
|