Websocket fix 1
This commit is contained in:
Urtzi Alfaro
@@ -5,19 +5,19 @@ WebSocket endpoints for real-time training progress updates
|
||||
|
||||
import json
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Dict, Any
|
||||
from fastapi import WebSocket, WebSocketDisconnect, Depends, HTTPException
|
||||
from fastapi.routing import APIRouter
|
||||
|
||||
import structlog
|
||||
logger = structlog.get_logger(__name__)
|
||||
|
||||
from app.services.messaging import training_publisher
|
||||
from shared.auth.decorators import (
|
||||
get_current_user_dep,
|
||||
get_current_tenant_id_dep
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Create WebSocket router
|
||||
websocket_router = APIRouter()
|
||||
|
||||
@@ -48,8 +48,9 @@ class ConnectionManager:
|
||||
logger.info(f"WebSocket disconnected for job {job_id}, connection {connection_id}")
|
||||
|
||||
async def send_to_job(self, job_id: str, message: dict):
|
||||
"""Send message to all connections for a specific job"""
|
||||
"""Send message to all connections for a specific job with better error handling"""
|
||||
if job_id not in self.active_connections:
|
||||
logger.debug(f"No active connections for job {job_id}")
|
||||
return
|
||||
|
||||
# Send to all connections for this job
|
||||
@@ -58,6 +59,7 @@ class ConnectionManager:
|
||||
for connection_id, websocket in self.active_connections[job_id].items():
|
||||
try:
|
||||
await websocket.send_json(message)
|
||||
logger.debug(f"📤 Sent {message.get('type', 'unknown')} to connection {connection_id}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to send message to connection {connection_id}: {e}")
|
||||
disconnected_connections.append(connection_id)
|
||||
@@ -65,6 +67,11 @@ class ConnectionManager:
|
||||
# Clean up disconnected connections
|
||||
for connection_id in disconnected_connections:
|
||||
self.disconnect(job_id, connection_id)
|
||||
|
||||
# Log successful sends
|
||||
active_count = len(self.active_connections.get(job_id, {}))
|
||||
if active_count > 0:
|
||||
logger.info(f"📡 Sent {message.get('type', 'unknown')} message to {active_count} connection(s) for job {job_id}")
|
||||
|
||||
# Global connection manager
|
||||
connection_manager = ConnectionManager()
|
||||
@@ -76,40 +83,31 @@ async def training_progress_websocket(
|
||||
job_id: str
|
||||
):
|
||||
"""
|
||||
WebSocket endpoint for real-time training progress updates
|
||||
|
||||
Message format sent to client:
|
||||
{
|
||||
"type": "progress" | "completed" | "failed" | "step_completed",
|
||||
"job_id": "job_123",
|
||||
"timestamp": "2025-07-30T19:08:53Z",
|
||||
"data": {
|
||||
"progress": 45,
|
||||
"current_step": "Training model for bread",
|
||||
"current_product": "bread",
|
||||
"products_completed": 2,
|
||||
"products_total": 5,
|
||||
"estimated_time_remaining_minutes": 8
|
||||
}
|
||||
}
|
||||
FIXED WebSocket endpoint for real-time training progress updates
|
||||
Prevents premature disconnection during training
|
||||
"""
|
||||
connection_id = f"{tenant_id}_{id(websocket)}"
|
||||
|
||||
# Accept connection
|
||||
await connection_manager.connect(websocket, job_id, connection_id)
|
||||
logger.info(f"WebSocket connection established for job {job_id}")
|
||||
|
||||
# Set up RabbitMQ consumer for this job
|
||||
consumer_task = None
|
||||
training_completed = False # Track training completion
|
||||
|
||||
try:
|
||||
# Start RabbitMQ consumer
|
||||
logger.info(f"Setting up RabbitMQ consumer for job {job_id}")
|
||||
consumer_task = asyncio.create_task(
|
||||
setup_rabbitmq_consumer_for_job(job_id, tenant_id)
|
||||
)
|
||||
|
||||
# Give consumer time to set up
|
||||
await asyncio.sleep(0.5)
|
||||
|
||||
# Send initial status if available
|
||||
try:
|
||||
# You can fetch current job status from database here
|
||||
initial_status = await get_current_job_status(job_id, tenant_id)
|
||||
if initial_status:
|
||||
await websocket.send_json({
|
||||
@@ -117,48 +115,92 @@ async def training_progress_websocket(
|
||||
"job_id": job_id,
|
||||
"data": initial_status
|
||||
})
|
||||
logger.info(f"Sent initial status for job {job_id}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to send initial status: {e}")
|
||||
|
||||
# Keep connection alive and handle client messages
|
||||
while True:
|
||||
# Keep connection alive - IMPROVED ERROR HANDLING
|
||||
last_activity = asyncio.get_event_loop().time()
|
||||
|
||||
while not training_completed:
|
||||
try:
|
||||
# Wait for client ping or other messages
|
||||
message = await websocket.receive_text()
|
||||
|
||||
if message == "ping":
|
||||
await websocket.send_text("pong")
|
||||
elif message == "get_status":
|
||||
# Send current status on demand
|
||||
current_status = await get_current_job_status(job_id, tenant_id)
|
||||
if current_status:
|
||||
# Wait for client messages with timeout
|
||||
try:
|
||||
message = await asyncio.wait_for(websocket.receive_text(), timeout=30.0)
|
||||
last_activity = asyncio.get_event_loop().time()
|
||||
|
||||
if message == "ping":
|
||||
await websocket.send_text("pong")
|
||||
logger.debug(f"Ping received from job {job_id}")
|
||||
elif message == "get_status":
|
||||
current_status = await get_current_job_status(job_id, tenant_id)
|
||||
if current_status:
|
||||
await websocket.send_json({
|
||||
"type": "current_status",
|
||||
"job_id": job_id,
|
||||
"data": current_status
|
||||
})
|
||||
elif message == "close":
|
||||
logger.info(f"Client requested connection close for job {job_id}")
|
||||
break
|
||||
|
||||
except asyncio.TimeoutError:
|
||||
# No message received in 30 seconds - send heartbeat
|
||||
current_time = asyncio.get_event_loop().time()
|
||||
if current_time - last_activity > 60: # 60 seconds of inactivity
|
||||
logger.warning(f"Long inactivity period for job {job_id}, sending heartbeat")
|
||||
|
||||
try:
|
||||
await websocket.send_json({
|
||||
"type": "current_status",
|
||||
"type": "heartbeat",
|
||||
"job_id": job_id,
|
||||
"data": current_status
|
||||
"timestamp": datetime.utcnow().isoformat()
|
||||
})
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to send heartbeat for job {job_id}: {e}")
|
||||
break
|
||||
|
||||
except WebSocketDisconnect:
|
||||
logger.info(f"WebSocket disconnected for job {job_id}")
|
||||
logger.info(f"WebSocket client disconnected for job {job_id}")
|
||||
break
|
||||
except Exception as e:
|
||||
logger.error(f"WebSocket error for job {job_id}: {e}")
|
||||
break
|
||||
# Don't break immediately - try to recover
|
||||
await asyncio.sleep(1)
|
||||
|
||||
logger.info(f"WebSocket loop ended for job {job_id}, training_completed: {training_completed}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Critical WebSocket error for job {job_id}: {e}")
|
||||
|
||||
finally:
|
||||
# Clean up
|
||||
# IMPROVED CLEANUP - Don't cancel consumer unless truly disconnecting
|
||||
logger.info(f"Cleaning up WebSocket connection for job {job_id}")
|
||||
connection_manager.disconnect(job_id, connection_id)
|
||||
|
||||
# Only cancel consumer if we're truly done (not just a temporary error)
|
||||
if consumer_task and not consumer_task.done():
|
||||
consumer_task.cancel()
|
||||
if training_completed:
|
||||
logger.info(f"Training completed, cancelling consumer for job {job_id}")
|
||||
consumer_task.cancel()
|
||||
else:
|
||||
logger.warning(f"WebSocket disconnected but training not completed for job {job_id}")
|
||||
# Let the consumer continue running for other potential connections
|
||||
# Don't cancel it unless we're sure the job is done
|
||||
|
||||
try:
|
||||
await consumer_task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
logger.info(f"Consumer task cancelled for job {job_id}")
|
||||
except Exception as e:
|
||||
logger.error(f"Consumer task error for job {job_id}: {e}")
|
||||
|
||||
|
||||
async def setup_rabbitmq_consumer_for_job(job_id: str, tenant_id: str):
|
||||
"""Set up RabbitMQ consumer to listen for training events for a specific job"""
|
||||
|
||||
logger.info(f"🚀 Setting up RabbitMQ consumer for job {job_id}")
|
||||
|
||||
try:
|
||||
# Create a unique queue for this WebSocket connection
|
||||
queue_name = f"websocket_training_{job_id}_{tenant_id}"
|
||||
@@ -170,12 +212,16 @@ async def setup_rabbitmq_consumer_for_job(job_id: str, tenant_id: str):
|
||||
body = message.body.decode()
|
||||
data = json.loads(body)
|
||||
|
||||
logger.debug(f"🔍 Received message for job {job_id}: {data.get('event_type', 'unknown')}")
|
||||
|
||||
# Extract event data
|
||||
event_type = data.get("event_type", "unknown")
|
||||
event_data = data.get("data", {})
|
||||
|
||||
# Only process messages for this specific job
|
||||
if event_data.get("job_id") != job_id:
|
||||
message_job_id = event_data.get("job_id") if event_data else None
|
||||
if message_job_id != job_id:
|
||||
logger.debug(f"⏭️ Ignoring message for different job: {message_job_id}")
|
||||
await message.ack()
|
||||
return
|
||||
|
||||
@@ -187,36 +233,76 @@ async def setup_rabbitmq_consumer_for_job(job_id: str, tenant_id: str):
|
||||
"data": event_data
|
||||
}
|
||||
|
||||
logger.info(f"📤 Forwarding {event_type} message to WebSocket clients for job {job_id}")
|
||||
|
||||
# Send to all WebSocket connections for this job
|
||||
await connection_manager.send_to_job(job_id, websocket_message)
|
||||
|
||||
# Check if this is a completion message
|
||||
if event_type in ["training.completed", "training.failed"]:
|
||||
logger.info(f"🎯 Training completion detected for job {job_id}: {event_type}")
|
||||
# Mark training as completed (you might want to store this in a global state)
|
||||
# For now, we'll let the WebSocket handle this through the message
|
||||
|
||||
# Acknowledge the message
|
||||
await message.ack()
|
||||
|
||||
logger.debug(f"Forwarded training event to WebSocket: {event_type}")
|
||||
logger.debug(f"✅ Successfully processed {event_type} for job {job_id}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error handling training message for WebSocket: {e}")
|
||||
logger.error(f"❌ Error handling training message for job {job_id}: {e}")
|
||||
import traceback
|
||||
logger.error(f"💥 Traceback: {traceback.format_exc()}")
|
||||
await message.nack(requeue=False)
|
||||
|
||||
# Check if training_publisher is connected
|
||||
if not training_publisher.connected:
|
||||
logger.warning(f"⚠️ Training publisher not connected for job {job_id}, attempting to connect...")
|
||||
success = await training_publisher.connect()
|
||||
if not success:
|
||||
logger.error(f"❌ Failed to connect training_publisher for job {job_id}")
|
||||
return
|
||||
|
||||
# Subscribe to training events
|
||||
await training_publisher.consume_events(
|
||||
logger.info(f"🔗 Subscribing to training events for job {job_id}")
|
||||
success = await training_publisher.consume_events(
|
||||
exchange_name="training.events",
|
||||
queue_name=queue_name,
|
||||
routing_key="training.*", # Listen to all training events
|
||||
callback=handle_training_message
|
||||
)
|
||||
|
||||
if success:
|
||||
logger.info(f"✅ Successfully set up RabbitMQ consumer for job {job_id} (queue: {queue_name})")
|
||||
|
||||
# Keep the consumer running indefinitely until cancelled
|
||||
try:
|
||||
while True:
|
||||
await asyncio.sleep(10) # Keep consumer alive
|
||||
logger.debug(f"🔄 Consumer heartbeat for job {job_id}")
|
||||
|
||||
except asyncio.CancelledError:
|
||||
logger.info(f"🛑 Consumer cancelled for job {job_id}")
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"💥 Consumer error for job {job_id}: {e}")
|
||||
raise
|
||||
else:
|
||||
logger.error(f"❌ Failed to set up RabbitMQ consumer for job {job_id}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to set up RabbitMQ consumer for WebSocket: {e}")
|
||||
logger.error(f"💥 Exception in setup_rabbitmq_consumer_for_job for job {job_id}: {e}")
|
||||
import traceback
|
||||
logger.error(f"🔥 Traceback: {traceback.format_exc()}")
|
||||
|
||||
|
||||
def map_event_type_to_websocket_type(rabbitmq_event_type: str) -> str:
|
||||
"""Map RabbitMQ event types to WebSocket message types"""
|
||||
mapping = {
|
||||
"training.started": "started",
|
||||
"training.progress": "progress",
|
||||
"training.completed": "completed",
|
||||
"training.failed": "failed",
|
||||
"training.completed": "completed", # This is the key completion event
|
||||
"training.failed": "failed", # This is also a completion event
|
||||
"training.cancelled": "cancelled",
|
||||
"training.step.completed": "step_completed",
|
||||
"training.product.started": "product_started",
|
||||
|
||||
Reference in New Issue
Block a user