Spaces:
Sleeping
Sleeping
| import WebSocket from 'ws'; | |
| import http from 'http'; | |
| import { createApp } from '../src/app'; | |
| import { WebSocketHandler } from '../src/websocket/wsHandler'; | |
| import { TaskBrokerWebSocketServer } from '../src/websocket/wsServer'; | |
| async function runTestSimulation() { | |
| console.log('--- Starting Distributed Task Broker Simulation Test ---'); | |
| // Setup app & server on test port 7899 | |
| const context = createApp(); | |
| const server = http.createServer(context.app); | |
| const wsHandler = new WebSocketHandler( | |
| context.workerService, | |
| context.taskQueueService, | |
| context.progressService, | |
| context.resultService | |
| ); | |
| const wsServer = new TaskBrokerWebSocketServer(server, wsHandler); | |
| context.adminController.setWsServer(wsServer); | |
| const PORT = 7899; | |
| await new Promise<void>((resolve) => server.listen(PORT, resolve)); | |
| console.log(`Test server running on port ${PORT}`); | |
| // Create initial batch tasks (100,000 rows, 10,000 chunkSize = 10 tasks) | |
| await context.taskQueueService.generateBatchTasks(100000, 10000); | |
| // Helper to create a simulated worker | |
| const simulateWorker = (id: string, name: string, platform: 'colab' | 'kaggle' | 'local', gpuName: string) => { | |
| return new Promise<void>((resolve, reject) => { | |
| const ws = new WebSocket(`ws://localhost:${PORT}/ws`); | |
| ws.on('open', () => { | |
| console.log(`[${id}] WebSocket connected. Registering worker...`); | |
| ws.send( | |
| JSON.stringify({ | |
| type: 'register', | |
| payload: { | |
| workerId: id, | |
| workerName: name, | |
| platform: platform, | |
| specs: { cpu: 8, ram: '32GB', gpu: gpuName, batchSize: 64 }, | |
| }, | |
| }) | |
| ); | |
| }); | |
| ws.on('message', (data: WebSocket.RawData) => { | |
| const msg = JSON.parse(data.toString()); | |
| // console.log(`[${id}] Received server message:`, msg.type); | |
| if (msg.type === 'registered') { | |
| console.log(`[${id}] Registered successfully! Requesting first OCR task...`); | |
| ws.send(JSON.stringify({ type: 'request_task', payload: { workerId: id } })); | |
| } else if (msg.type === 'task_assignment') { | |
| const task = msg.payload; | |
| console.log( | |
| `[${id}] ASSIGNED TASK: ${task.taskId} | Rows: ${task.startRow} -> ${task.endRow} (${task.totalRows} rows) | Offsets: ${task.startOffset} -> ${task.endOffset}` | |
| ); | |
| // Send heartbeat and progress update | |
| ws.send( | |
| JSON.stringify({ | |
| type: 'progress', | |
| payload: { workerId: id, taskId: task.taskId, progress: 50, rowsProcessed: task.totalRows / 2, currentSpeed: 250 }, | |
| }) | |
| ); | |
| // Complete task after short delay | |
| setTimeout(() => { | |
| console.log(`[${id}] Completing task ${task.taskId}...`); | |
| ws.send( | |
| JSON.stringify({ | |
| type: 'task_completed', | |
| payload: { workerId: id, taskId: task.taskId, rowsProcessed: task.totalRows }, | |
| }) | |
| ); | |
| }, 300); | |
| } else if (msg.type === 'wait') { | |
| console.log(`[${id}] Queue empty or wait state received: ${msg.payload.reason}`); | |
| ws.close(); | |
| resolve(); | |
| } | |
| }); | |
| ws.on('error', (err) => { | |
| console.error(`[${id}] WS Error:`, err); | |
| reject(err); | |
| }); | |
| }); | |
| }; | |
| // Run 3 concurrent remote workers | |
| await Promise.all([ | |
| simulateWorker('colab-worker-1', 'Google Colab T4 Node', 'colab', 'NVIDIA T4'), | |
| simulateWorker('kaggle-worker-1', 'Kaggle P100 Node', 'kaggle', 'NVIDIA P100'), | |
| simulateWorker('local-worker-1', 'Desktop Workstation', 'local', 'RTX 4090'), | |
| ]); | |
| // Check system statistics after simulation | |
| const stats = await context.statisticsService.getSystemStatistics(); | |
| console.log('\n--- Final System Statistics ---'); | |
| console.log(`Total Tasks: ${stats.tasks.total}`); | |
| console.log(`Completed Tasks: ${stats.tasks.completed}`); | |
| console.log(`Pending Tasks: ${stats.tasks.pending}`); | |
| console.log(`Total Rows Processed: ${stats.throughput.totalRowsProcessed}`); | |
| console.log(`Completion Percentage: ${stats.tasks.completionPercentage}%`); | |
| server.close(); | |
| console.log('--- Simulation Completed Successfully ---'); | |
| } | |
| runTestSimulation().catch(console.error); | |