Spaces:
Sleeping
Sleeping
| import crypto from 'crypto'; | |
| import { IWorkerService } from '../interfaces/IWorkerService'; | |
| import { IWorkerRepository } from '../../repositories/interfaces/IWorkerRepository'; | |
| import { IWorker, RegisterWorkerInput, WorkerStatus } from '../../types/worker.types'; | |
| import { WorkerModel } from '../../models/worker.model'; | |
| import { CONFIG } from '../../config/app.config'; | |
| import { logger } from '../../utils/logger'; | |
| export class WorkerService implements IWorkerService { | |
| private workerRepo: IWorkerRepository; | |
| constructor(workerRepo: IWorkerRepository) { | |
| this.workerRepo = workerRepo; | |
| } | |
| public async registerWorker(input: RegisterWorkerInput): Promise<IWorker> { | |
| const workerId = input.workerId || `worker-${Math.random().toString(36).substring(2, 9)}`; | |
| const existing = await this.workerRepo.findById(workerId); | |
| const specs = input.specs || {}; | |
| const cpu = specs.cpu !== undefined ? String(specs.cpu) : String(input.cpu || 'Unknown'); | |
| const ram = specs.ram !== undefined ? String(specs.ram) : String(input.ram || 'Unknown'); | |
| const gpu = specs.gpu !== undefined ? specs.gpu : input.gpu || 'None'; | |
| const tpu = specs.tpu !== undefined ? specs.tpu : input.tpu || 'None'; | |
| const batchSize = specs.batchSize || input.batchSize || 32; | |
| const workerModel = new WorkerModel({ | |
| ...(existing || {}), | |
| workerId, | |
| workerName: input.workerName, | |
| platform: input.platform, | |
| cpu, | |
| ram, | |
| gpu, | |
| tpu, | |
| batchSize, | |
| lastHeartbeat: new Date().toISOString(), | |
| status: 'idle', | |
| currentTaskId: existing?.currentTaskId || null, | |
| }); | |
| const saved = await this.workerRepo.save(workerModel); | |
| logger.info(`Registered worker [${saved.workerId}] (${saved.workerName}) on platform [${saved.platform}]`); | |
| return saved; | |
| } | |
| public async getWorker(workerId: string): Promise<IWorker | null> { | |
| return this.workerRepo.findById(workerId); | |
| } | |
| public async getAllWorkers(status?: WorkerStatus): Promise<IWorker[]> { | |
| return this.workerRepo.findAll(status); | |
| } | |
| public async recordHeartbeat( | |
| workerId: string, | |
| currentTaskId?: string, | |
| rowsProcessed?: number, | |
| currentSpeed?: number | |
| ): Promise<IWorker | null> { | |
| const worker = await this.workerRepo.findById(workerId); | |
| if (!worker) return null; | |
| worker.lastHeartbeat = new Date().toISOString(); | |
| if (currentSpeed !== undefined && currentSpeed > 0) { | |
| // Exponential moving average for smooth throughput tracking | |
| worker.rowsPerSecond = worker.rowsPerSecond > 0 | |
| ? Math.round(worker.rowsPerSecond * 0.7 + currentSpeed * 0.3) | |
| : Math.round(currentSpeed); | |
| } | |
| if (currentTaskId !== undefined) { | |
| worker.currentTaskId = currentTaskId; | |
| worker.status = currentTaskId ? 'busy' : 'idle'; | |
| } | |
| return this.workerRepo.save(worker); | |
| } | |
| public async setWorkerStatus( | |
| workerId: string, | |
| status: WorkerStatus, | |
| currentTaskId?: string | null | |
| ): Promise<IWorker | null> { | |
| const worker = await this.workerRepo.findById(workerId); | |
| if (!worker) return null; | |
| worker.status = status; | |
| if (currentTaskId !== undefined) { | |
| worker.currentTaskId = currentTaskId; | |
| } | |
| worker.lastHeartbeat = new Date().toISOString(); | |
| return this.workerRepo.save(worker); | |
| } | |
| public async updateWorkerStats( | |
| workerId: string, | |
| completedIncrement: number, | |
| failedIncrement: number, | |
| rowsProcessedIncrement: number | |
| ): Promise<void> { | |
| const worker = await this.workerRepo.findById(workerId); | |
| if (!worker) return; | |
| worker.tasksCompleted += completedIncrement; | |
| worker.tasksFailed += failedIncrement; | |
| worker.totalRowsProcessed += rowsProcessedIncrement; | |
| await this.workerRepo.save(worker); | |
| } | |
| public async markStaleWorkersOffline(): Promise<IWorker[]> { | |
| const thresholdMs = Date.now() - (CONFIG.LEASE.HEARTBEAT_INTERVAL_SECONDS * 2.5 * 1000); | |
| const thresholdIso = new Date(thresholdMs).toISOString(); | |
| const staleWorkers = await this.workerRepo.findStaleWorkers(thresholdIso); | |
| const markedOffline: IWorker[] = []; | |
| for (const w of staleWorkers) { | |
| logger.warn(`Worker [${w.workerId}] missed heartbeats. Marking offline.`); | |
| w.status = 'offline'; | |
| w.currentTaskId = null; | |
| const updated = await this.workerRepo.save(w); | |
| markedOffline.push(updated); | |
| } | |
| return markedOffline; | |
| } | |
| public async removeWorker(workerId: string): Promise<boolean> { | |
| return this.workerRepo.delete(workerId); | |
| } | |
| } | |