160 lines · 5.4 KB · Python
RawDownload
"""Server-Sent Events manager for real-time updates. This module provides a centralized event manager for broadcastingreal-time updates to connected clients via SSE.""" import asyncioimport jsonimport loggingfrom collections import defaultdictfrom dataclasses import dataclass, fieldfrom datetime import datetimefrom typing import Any, AsyncIterator, Dict, List, Optional, Set logger = logging.getLogger(__name__)  @dataclassclass Event:    """SSE Event with type, data, and timestamp."""     type: str    data: Dict[str, Any]    timestamp: str = field(default_factory=lambda: datetime.utcnow().isoformat())    id: str = field(default_factory=lambda: str(int(datetime.utcnow().timestamp() * 1000)))     def to_sse(self) -> str:        """Convert to SSE wire format.         SSE format:        id: <event_id>        event: <event_type>        data: <json_payload>         (blank line to delimit event)        """        payload = {"data": self.data, "timestamp": self.timestamp}        lines = [f"id: {self.id}", f"event: {self.type}", f"data: {json.dumps(payload)}", "", ""]        return "\n".join(lines)  class EventManager:    """Manages SSE connections and broadcasts events to subscribers.     Supports multiple channels for targeted broadcasting.    Thread-safe for async operations.    """     def __init__(self):        self._subscribers: Dict[str, Set[asyncio.Queue]] = defaultdict(set)        self._lock = asyncio.Lock()        self._event_count = 0     async def subscribe(self, channel: str = "default") -> AsyncIterator[Event]:        """Subscribe to events on a specific channel.         Args:            channel: Channel name to subscribe to (default: "default")         Yields:            Event objects as they're published to the channel        """        queue: asyncio.Queue[Event] = asyncio.Queue(maxsize=100)         async with self._lock:            self._subscribers[channel].add(queue)            logger.info(                f"New subscriber to channel '{channel}' (total: {len(self._subscribers[channel])})"            )         try:            while True:                event = await queue.get()                yield event        except asyncio.CancelledError:            logger.debug(f"Subscriber cancelled on channel '{channel}'")        except Exception as e:            logger.error(f"Error in subscriber loop: {e}")        finally:            async with self._lock:                self._subscribers[channel].discard(queue)                logger.info(                    f"Subscriber left channel '{channel}' (remaining: {len(self._subscribers[channel])})"                )     async def publish(self, event: Event, channel: str = "default"):        """Publish event to all subscribers on a channel.         Handles queue full errors gracefully by dropping events.         Args:            event: Event to publish            channel: Channel to publish to (default: "default")        """        async with self._lock:            subscribers = self._subscribers[channel]            if not subscribers:                return  # No subscribers, skip             self._event_count += 1            dead_queues = set()             for queue in subscribers:                try:                    queue.put_nowait(event)                except asyncio.QueueFull:                    logger.warning(                        f"Queue full for channel '{channel}', dropping event (type: {event.type})"                    )                    dead_queues.add(queue)                except Exception as e:                    logger.error(f"Error publishing to subscriber: {e}")                    dead_queues.add(queue)             # Clean up dead queues            self._subscribers[channel] -= dead_queues     async def publish_status(self, status_data: Dict[str, Any]):        """Publish status update event."""        await self.publish(Event(type="status", data=status_data))     async def publish_gpu(self, gpu_data: List[Dict[str, Any]]):        """Publish GPU metrics event."""        await self.publish(Event(type="gpu", data={"gpus": gpu_data, "count": len(gpu_data)}))     async def publish_metrics(self, metrics_data: Dict[str, Any]):        """Publish vLLM performance metrics event."""        await self.publish(Event(type="metrics", data=metrics_data))     async def publish_log_line(self, session_id: str, line: str):        """Publish individual log line to session-specific channel."""        await self.publish(            Event(type="log", data={"session_id": session_id, "line": line}),            channel=f"logs:{session_id}",        )     async def publish_launch_progress(        self, recipe_id: str, stage: str, message: str, progress: Optional[float] = None    ):        """Publish model launch progress event.         Stages: evicting, launching, waiting, ready, error        """        data = {"recipe_id": recipe_id, "stage": stage, "message": message}        if progress is not None:            data["progress"] = progress         await self.publish(Event(type="launch_progress", data=data))     def get_stats(self) -> Dict[str, Any]:        """Get event manager statistics."""        return {            "total_events_published": self._event_count,            "channels": {channel: len(subs) for channel, subs in self._subscribers.items()},            "total_subscribers": sum(len(subs) for subs in self._subscribers.values()),        }  # Global singleton event managerevent_manager = EventManager()