94 lines
2.9 KiB
Python
94 lines
2.9 KiB
Python
"""Data broadcaster for WebSocket connections"""
|
|
|
|
import asyncio
|
|
from datetime import datetime
|
|
from typing import Dict, Any, Optional
|
|
|
|
from app.core.websocket.manager import manager
|
|
|
|
|
|
class DataBroadcaster:
|
|
"""Periodically broadcasts data to connected WebSocket clients"""
|
|
|
|
def __init__(self):
|
|
self.running = False
|
|
self.tasks: Dict[str, asyncio.Task] = {}
|
|
|
|
async def get_dashboard_stats(self) -> Dict[str, Any]:
|
|
"""Get dashboard statistics"""
|
|
return {
|
|
"total_datasources": 9,
|
|
"active_datasources": 8,
|
|
"tasks_today": 45,
|
|
"success_rate": 97.8,
|
|
"last_updated": datetime.utcnow().isoformat(),
|
|
"alerts": {"critical": 0, "warning": 2, "info": 5},
|
|
}
|
|
|
|
async def broadcast_stats(self, interval: int = 5):
|
|
"""Broadcast dashboard stats periodically"""
|
|
while self.running:
|
|
try:
|
|
stats = await self.get_dashboard_stats()
|
|
await manager.broadcast(
|
|
{
|
|
"type": "data_frame",
|
|
"channel": "dashboard",
|
|
"timestamp": datetime.utcnow().isoformat(),
|
|
"payload": {"stats": stats},
|
|
},
|
|
channel="dashboard",
|
|
)
|
|
except Exception:
|
|
pass
|
|
await asyncio.sleep(interval)
|
|
|
|
async def broadcast_alert(self, alert: Dict[str, Any]):
|
|
"""Broadcast an alert to all connected clients"""
|
|
await manager.broadcast(
|
|
{
|
|
"type": "alert_notification",
|
|
"timestamp": datetime.utcnow().isoformat(),
|
|
"data": {"alert": alert},
|
|
}
|
|
)
|
|
|
|
async def broadcast_gpu_update(self, data: Dict[str, Any]):
|
|
"""Broadcast GPU cluster update"""
|
|
await manager.broadcast(
|
|
{
|
|
"type": "data_frame",
|
|
"channel": "gpu_clusters",
|
|
"timestamp": datetime.utcnow().isoformat(),
|
|
"payload": data,
|
|
}
|
|
)
|
|
|
|
async def broadcast_custom(self, channel: str, data: Dict[str, Any]):
|
|
"""Broadcast custom data to a specific channel"""
|
|
await manager.broadcast(
|
|
{
|
|
"type": "data_frame",
|
|
"channel": channel,
|
|
"timestamp": datetime.utcnow().isoformat(),
|
|
"payload": data,
|
|
},
|
|
channel=channel if channel in manager.active_connections else "all",
|
|
)
|
|
|
|
def start(self):
|
|
"""Start all broadcasters"""
|
|
if not self.running:
|
|
self.running = True
|
|
self.tasks["dashboard"] = asyncio.create_task(self.broadcast_stats(5))
|
|
|
|
def stop(self):
|
|
"""Stop all broadcasters"""
|
|
self.running = False
|
|
for task in self.tasks.values():
|
|
task.cancel()
|
|
self.tasks.clear()
|
|
|
|
|
|
broadcaster = DataBroadcaster()
|