WebSocket Real-Time Updates Implementation¶

Status: ✅ Implemented (Phase 1-4 Complete) Date: 2025-01-14

Overview¶

This implementation replaces the polling-based approach with WebSocket real-time updates for task and experiment monitoring. The system uses FastAPI WebSockets on the backend and React hooks on the frontend to provide instant updates without continuous API polling.

Architecture¶

Backend Components¶

1. Event Broadcaster (`src/web/events/broadcaster.py`)¶

Purpose: In-memory pub/sub system for event distribution to WebSocket clients

Key Features:

asyncio.Queue-based event distribution
Thread-safe with asyncio locks
Support for both async (FastAPI) and sync (ARQ worker) contexts
Automatic queue overflow handling (drops oldest events)
Subscriber management per task ID

Core Classes:

class EventBroadcaster:
    async def subscribe(task_id) -> asyncio.Queue
    async def unsubscribe(task_id, queue)
    async def broadcast(task_id, event)  # Async version
    def broadcast_sync(task_id, event)   # Sync wrapper for ARQ workers

Event Types:

TASK_STARTED - Task execution begins
TASK_PROGRESS - Periodic progress updates
TASK_COMPLETED / TASK_FAILED - Task completion
EXPERIMENT_STARTED - Experiment deployment starts
EXPERIMENT_PROGRESS - Benchmark phase starts
EXPERIMENT_COMPLETED / EXPERIMENT_FAILED - Experiment finishes
BENCHMARK_STARTED / BENCHMARK_PROGRESS - Benchmark-specific events

2. WebSocket Routes (`src/web/routes/websocket.py`)¶

Endpoints:

/api/ws/tasks/{task_id} - Main WebSocket endpoint for task updates

Accepts WebSocket connections from clients
Subscribes to EventBroadcaster for the specified task
Sends JSON events to client in real-time
Handles disconnections gracefully

/api/ws/experiments/{experiment_id} - Experiment-specific updates

For fine-grained experiment monitoring
Uses unique subscription keys to avoid collisions

GET /api/ws/tasks/{task_id}/subscribers - Monitoring endpoint

Returns current number of active WebSocket connections
Useful for debugging

3. ARQ Worker Integration (`src/web/workers/autotuner_worker.py`)¶

Event Broadcasting Points:

Task Started (Line 184-191)

broadcaster.broadcast_sync(task_id, create_event(
    EventType.TASK_STARTED, task_id=task_id, ...
))

Experiment Started (Line 334-346)
- Broadcasts when experiment deployment begins
- Includes parameters and initial status
Benchmark Progress (Line 365-374)
- Broadcasts when benchmark phase begins
- Updated through async monitor task
Experiment Completed (Line 453-467)
- Broadcasts on experiment success/failure
- Includes metrics, objective score, elapsed time
Task Progress (Line 484-498)
- Broadcasts after each experiment completes
- Includes progress percentage and best score
Task Completed (Line 596-611)
- Broadcasts on task completion
- Includes final summary statistics

Frontend Components¶

1. Base WebSocket Hook (`frontend/src/hooks/useWebSocket.ts`)¶

Purpose: Low-level WebSocket connection management with reconnection logic

Features:

Automatic Reconnection: Exponential backoff with jitter
Configurable Parameters:
- reconnectDelay: Base delay (default 1000ms)
- maxReconnectDelay: Maximum delay cap (default 30000ms)
- maxReconnectAttempts: Retry limit (default: Infinity)
Message History: Keeps last 100 messages
Connection State: CONNECTING, OPEN, CLOSING, CLOSED
Clean Lifecycle Management: Proper cleanup on unmount

API:

const {
  state,              // Current connection state
  lastMessage,        // Latest message received
  messageHistory,     // Last 100 messages
  sendMessage,        // Send message to server
  connect,            // Manual connect
  disconnect,         // Manual disconnect
  isConnected,        // Boolean convenience
  reconnectAttempts   // Number of retry attempts
} = useWebSocket(url, options);

Reconnection Algorithm:

// Exponential backoff with jitter
delay = min(baseDelay * 2^attempt, maxDelay)
jitter = delay * 0.25 * random(-1, 1)
finalDelay = delay + jitter

2. Task-Specific Hook (`frontend/src/hooks/useTaskWebSocket.ts`)¶

Purpose: High-level hook for task monitoring with React Query integration

Features:

Automatically constructs WebSocket URL from task ID
Invalidates React Query caches on events
Type-safe event handling
Console logging for debugging

Event to Cache Mapping:

task_* events → Invalidate ["tasks"] and ["task", taskId]
experiment_* events → Invalidate ["experiments", taskId]
benchmark_* events → Invalidate ["experiments", taskId]

Usage:

// Automatically connects if taskId is not null
useTaskWebSocket(taskId, enabled);

3. Page Integration¶

Tasks Page (frontend/src/pages/Tasks.tsx):

Finds first running task from task list
Connects WebSocket only for running tasks
Reduced polling interval from 5s to 30s (fallback only)
WebSocket provides real-time updates

Experiments Page (frontend/src/pages/Experiments.tsx):

Similar pattern to Tasks page
Connects to running task’s WebSocket
Automatically updates experiment list

Communication Flow¶

┌─────────────────┐
│  ARQ Worker     │
│  (Background)   │
└────────┬────────┘
         │
         │ broadcast_sync()
         ↓
┌─────────────────┐
│ EventBroadcaster│ (In-memory pub/sub)
│  Global Instance│
└────────┬────────┘
         │
         │ asyncio.Queue
         ↓
┌─────────────────┐
│ WebSocket Route │ /api/ws/tasks/{id}
│  (FastAPI)      │
└────────┬────────┘
         │
         │ WebSocket Protocol
         ↓
┌─────────────────┐
│ useWebSocket    │ (React Hook)
│  Frontend       │
└────────┬────────┘
         │
         │ Event Callback
         ↓
┌─────────────────┐
│useTaskWebSocket │ (React Hook)
│                 │
└────────┬────────┘
         │
         │ queryClient.invalidateQueries()
         ↓
┌─────────────────┐
│ React Query     │ (Automatic refetch)
│  Cache          │
└────────┬────────┘
         │
         │ Component Re-render
         ↓
┌─────────────────┐
│   UI Update     │
└─────────────────┘

Benefits¶

Performance Improvements¶

Reduced Network Traffic
- Before: 1 API call every 5 seconds per page = 12 calls/minute
- After: 1 API call every 30 seconds + WebSocket events
- Reduction: ~83% fewer HTTP requests
Lower Latency
- Before: Average 2.5s delay (half of polling interval)
- After: < 100ms delay (WebSocket event propagation)
- Improvement: ~25x faster updates
Server Load
- Reduced database queries (fewer GET requests)
- Single WebSocket connection vs multiple HTTP requests
- More efficient for multiple concurrent users

User Experience Improvements¶

Real-Time Feedback
- Instant status updates when tasks start/complete
- Live experiment progress without page refresh
- Immediate error notifications
Progress Tracking
- Accurate progress percentage updates
- Current experiment number in real-time
- Best score updates as experiments complete
Reliability
- Automatic reconnection on network issues
- Fallback polling if WebSocket fails
- No stuck UI states
Optimized Polling Strategy
- Tasks Page: Adaptive polling based on task state
  - Running tasks: 30s fallback polling
  - Idle (no running tasks): 5min fallback polling
- ExperimentProgressBar: No polling, relies on parent WebSocket updates
- Experiments Page: No polling, WebSocket-driven updates only
- Dashboard: Intentional 5-10s polling for real-time system monitoring (GPU, Worker status)

Polling Optimization Details¶

Before WebSocket Implementation:

Tasks page: 5s unconditional polling
ExperimentProgressBar: 5s independent polling per task row
Total requests/minute: ~24 (2 endpoints × 12 polls)

After WebSocket Implementation:

Tasks page: 30s (running) / 5min (idle) adaptive polling
ExperimentProgressBar: No polling (WebSocket invalidation)
Experiments page: No polling (WebSocket invalidation)
WebSocket: Real-time events + targeted HTTP fetches
Total requests/minute (idle): ~0.4 (95% reduction)
Total requests/minute (running): ~4 (83% reduction)

Configuration¶

Backend Settings¶

No additional configuration required. WebSocket support is enabled automatically when FastAPI app starts.

Frontend Settings¶

WebSocket URL is automatically constructed:

const wsUrl = `${protocol === "https:" ? "wss:" : "ws:"}//${host}/api/ws/tasks/${taskId}`;

Reconnection Tuning¶

Modify useTaskWebSocket.ts parameters:

reconnectDelay: 1000,        // Initial retry delay (1s)
maxReconnectDelay: 10000,    // Maximum retry delay (10s)
maxReconnectAttempts: 10,    // Stop after 10 attempts

Debugging¶

Backend Logs¶

# Watch WebSocket events in worker log
tail -f logs/worker.log | grep "EventBroadcaster"

# Check WebSocket connections
curl http://localhost:8000/api/ws/tasks/1/subscribers

Frontend Logs¶

Open browser console to see:

[useTaskWebSocket] Connected to task 1
[useTaskWebSocket] Received event: {type: "task_started", ...}
[useWebSocket] Reconnecting in 1000ms (attempt 1/10)...

Verifying Polling Optimization¶

Method 1: Browser DevTools Network Tab

Open browser DevTools (F12)
Go to Network tab
Filter by “tasks” or “experiments”
Observe request intervals:
- Idle state (no running tasks): Should see requests ~5 minutes apart
- Running state (task executing): Should see requests ~30 seconds apart
- With WebSocket: Should see burst of requests only after WebSocket events

Method 2: Console Timing

// Paste in browser console to monitor request frequency
let lastRequest = {};
const originalFetch = window.fetch;
window.fetch = function(...args) {
  const url = args[0];
  if (url.includes('/api/tasks') || url.includes('/api/experiments')) {
    const now = Date.now();
    const last = lastRequest[url] || 0;
    const delta = now - last;
    console.log(`[Fetch Monitor] ${url} - ${delta}ms since last (${(delta/1000).toFixed(1)}s)`);
    lastRequest[url] = now;
  }
  return originalFetch.apply(this, args);
};

Expected Output (Idle):

[Fetch Monitor] /api/tasks - 300000ms since last (300.0s)  // 5 minutes
[Fetch Monitor] /api/experiments/task/1 - 0ms since last (0.0s)  // Triggered by WebSocket event

Expected Output (Running):

[Fetch Monitor] /api/tasks - 30000ms since last (30.0s)  // 30 seconds
[Fetch Monitor] /api/experiments/task/1 - 2500ms since last (2.5s)  // WebSocket event

Event Structure¶

{
  "type": "experiment_completed",
  "task_id": 1,
  "experiment_id": 3,
  "timestamp": 1705234567.89,
  "message": "Experiment 3 success",
  "data": {
    "status": "success",
    "metrics": {...},
    "objective_score": 2.45,
    "elapsed_time": 123.45
  }
}

Known Limitations¶

In-Memory Only: Events are not persisted. If server restarts, active WebSocket connections are lost.
Single-Server: Broadcasting only works within a single server instance (no Redis pub/sub yet).
Browser Tab Close: WebSocket disconnects immediately, no automatic reconnection from different tab.

Future Enhancements¶

Redis Pub/Sub: For multi-server deployments
Event History: Store last N events for late-joining clients
Bandwidth Optimization: Compress events, throttle high-frequency updates
Authentication: Add JWT token validation for WebSocket connections
Metrics: Track WebSocket connection count, event rates, latency

Code Statistics¶

Backend:

broadcaster.py: 238 lines
websocket.py (routes): 152 lines
autotuner_worker.py: +87 lines (event broadcasting)
Total Backend: ~477 lines

Frontend:

useWebSocket.ts: 302 lines
useTaskWebSocket.ts: 81 lines
Tasks.tsx: +13 lines (WebSocket integration)
Experiments.tsx: +13 lines (WebSocket integration)
Total Frontend: ~409 lines

Grand Total: ~886 lines of new code

Migration Notes¶

Backward Compatibility¶

The implementation is fully backward compatible:

Polling still works as fallback (30s interval)
WebSocket is additive, not replacing existing API
If WebSocket fails, UI continues to function with polling

Rollback Plan¶

If issues occur:

Remove WebSocket hook calls from Tasks.tsx and Experiments.tsx
Revert polling interval from 30s back to 5s
Restart frontend: npm run dev

No backend changes needed for rollback.

References¶

FastAPI WebSocket docs: https://fastapi.tiangolo.com/advanced/websockets/
React Query invalidation: https://tanstack.com/query/latest/docs/react/guides/query-invalidation
WebSocket reconnection patterns: https://javascript.info/websocket#reconnection