""" Motor del Worker, conexión WebSocket y ejecución de tareas de compresión. 100% idéntico a la lógica y ciclo de vida de hf_worker/main.go. """ import os import time import random import asyncio import threading import subprocess from pathlib import Path from typing import Optional, Dict, Any from urllib.parse import urlparse import psutil import requests import websockets from config import ( MANAGER_URL, WORKER_SECRET, WORKER_ID, WORK_DIR, MAX_JOBS, SPACE_URL, logger, ) from state import STATE, CurrentJob from protocol import ( TYPE_HANDSHAKE, TYPE_HANDSHAKE_ACK, TYPE_PING, TYPE_PONG, TYPE_TASK_ASSIGN, TYPE_TASK_PROGRESS, TYPE_TASK_COMPLETE, TYPE_TASK_FAILED, TYPE_TASK_CANCEL, build_ws_packet, parse_ws_packet, ) from ffmpeg import ( get_duration, compress_with_cmd, ) class UploadProgressReader: """Itera el archivo de salida y reporta velocidad/progreso al panel.""" def __init__(self, file_obj, total_size: int, task_id: str, worker): self.file_obj = file_obj self.total_size = total_size self.task_id = task_id self.worker = worker def __len__(self): return self.total_size def __iter__(self): started = time.monotonic() last_report = started uploaded = 0 while True: chunk = self.file_obj.read(1024 * 1024) if not chunk: break uploaded += len(chunk) now = time.monotonic() if now - last_report >= 1 or uploaded >= self.total_size: elapsed = max(now - started, 0.001) eta = "" if uploaded < self.total_size: eta = f"{int((self.total_size - uploaded) / (uploaded / elapsed))}s" progress = { "task_id": self.task_id, "stage": "uploading", "percent": min(100.0, uploaded * 100 / max(self.total_size, 1)), "current": uploaded, "total": self.total_size, "speed": f"{uploaded / elapsed / (1024 * 1024):.1f} MB/s", "eta": eta, } asyncio.run_coroutine_threadsafe( self.worker._send_ws_safe(TYPE_TASK_PROGRESS, progress), self.worker.loop, ) last_report = now yield chunk class WorkerEngine: """Cliente WebSocket y gestor de tareas del clúster.""" def __init__(self): self.manager_url: str = MANAGER_URL self.secret: str = WORKER_SECRET self.worker_id: str = WORKER_ID self.work_dir: str = WORK_DIR self.ws_conn: Optional[websockets.WebSocketClientProtocol] = None self.loop: Optional[asyncio.AbstractEventLoop] = None self.session = requests.Session() self.session.headers.update({ "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36", }) def start(self): """Inicia el bucle de supervisión en un hilo secundario continuo.""" thread = threading.Thread(target=self._run_event_loop, daemon=True, name="WorkerEngineThread") thread.start() def _run_event_loop(self): self.loop = asyncio.new_event_loop() asyncio.set_event_loop(self.loop) self.loop.run_until_complete(self._supervise_connection()) async def _supervise_connection(self): cores = psutil.cpu_count(logical=True) or 2 STATE.log(f"🖥️ Nodo de Inferencia: {self.worker_id} (Hilos CPU: {cores})") if SPACE_URL: STATE.log(f"💓 URL Watchdog detectada: {SPACE_URL}") STATE.log(f"📁 Directorio temporal: {self.work_dir}") STATE.log(f"🌐 Upstream Manager: {self.manager_url}") # Bucle continuo de conexión con backoff exponencial + jitter (idéntico a main.go) backoff = 2.0 while True: try: await self._run() if STATE.was_connected: backoff = 2.0 elif backoff < 60.0: backoff = min(60.0, backoff * 2.0) except Exception as e: STATE.set_connection(False, str(e)) STATE.log(f"⚠️ Desconectado del Manager: {e}. Reintentando en {int(round(backoff))}s...", level="WARNING") if STATE.was_connected: backoff = 2.0 elif backoff < 60.0: backoff = min(60.0, backoff * 2.0) jitter = random.uniform(0.0, 1.0) await asyncio.sleep(backoff + jitter) def _resolve_url(self, target_url: str) -> str: """Resuelve URLs relativas contra la URL base del Manager.""" if target_url.startswith("http://") or target_url.startswith("https://"): return target_url base = self.manager_url.rstrip("/") path = target_url.lstrip("/") return f"{base}/{path}" async def _send_ws_safe(self, msg_type: str, payload_obj: Any = None) -> bool: """Envía un paquete WS de forma segura.""" if self.ws_conn and self.ws_conn.open: try: await self.ws_conn.send(build_ws_packet(msg_type, payload_obj)) return True except Exception as e: logger.warning(f"Error enviando {msg_type}: {e}") return False return False async def _run(self): # Cada intento arranca marcado como no-conectado STATE.was_connected = False u = urlparse(self.manager_url) if not u.netloc and not u.path: raise ValueError(f"URL del manager inválida: {self.manager_url}") host = u.netloc or u.path ws_scheme = "wss" if u.scheme in ("https", "wss") else "ws" ws_url = f"{ws_scheme}://{host}/ws/worker" STATE.log(f"Intentando conectar a {ws_url}...") dial_headers = { "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36", } async with websockets.connect( ws_url, extra_headers=dial_headers, ping_interval=None, # Manejado manualmente por protocolo max_size=10 * 1024 * 1024, ) as ws: self.ws_conn = ws # 1. Enviar Handshake (incluyendo cualquier tarea activa si reconecta tras reinicio del Manager) active_tasks = STATE.get_active_tasks_for_handshake() cores = psutil.cpu_count(logical=True) or 2 handshake_payload = { "worker_id": self.worker_id, "secret": self.secret, "cpu_cores": cores, "max_jobs": MAX_JOBS, "space_url": SPACE_URL, "active_tasks": active_tasks, } await ws.send(build_ws_packet(TYPE_HANDSHAKE, handshake_payload)) # 2. Esperar HandshakeAck ack_raw = await ws.recv() mtype, ack_payload = parse_ws_packet(ack_raw) if mtype != TYPE_HANDSHAKE_ACK or not isinstance(ack_payload, dict) or ack_payload.get("status") != "ok": err_msg = ack_payload.get("message", "rechazado por manager") if isinstance(ack_payload, dict) else "rechazado" STATE.set_connection(False, err_msg) raise ConnectionRefusedError(f"handshake rechazado por manager: {err_msg}") STATE.set_connection(True) STATE.log("✅ ¡Conectado y autenticado con éxito en el Manager! Esperando tareas...") # 3. Tarea asíncrona de Ping periódico (cada 20s) ping_stop = asyncio.Event() async def ping_loop(): while not ping_stop.is_set(): try: await asyncio.sleep(20) memory = psutil.virtual_memory() try: disk = psutil.disk_usage('/') except Exception: disk = None telemetry = { "cpu_percent": psutil.cpu_percent(interval=None), "memory_percent": memory.percent, "memory_used_mb": int(memory.used / (1024 * 1024)), "memory_total_mb": int(memory.total / (1024 * 1024)), "disk_percent": disk.percent if disk else 0.0, "disk_used_mb": int(disk.used / (1024 * 1024)) if disk else 0, "disk_total_mb": int(disk.total / (1024 * 1024)) if disk else 0, } if not await self._send_ws_safe(TYPE_PING, telemetry): break except Exception: break ping_task = asyncio.create_task(ping_loop()) try: # 4. Bucle principal de lectura de mensajes async for raw_msg in ws: msg_type, payload = parse_ws_packet(raw_msg) if not msg_type: continue if msg_type == TYPE_PING: await self._send_ws_safe(TYPE_PONG, None) elif msg_type == TYPE_TASK_ASSIGN and isinstance(payload, dict): task_id = payload.get("task_id", "") file_name = payload.get("file_name", "") upload_url = payload.get("upload_url", "") job = CurrentJob(task_id, file_name, upload_url) STATE.add_job(job) # Ejecutar tarea en hilo independiente asyncio.create_task(self._execute_task_wrapper(payload, job)) elif msg_type == TYPE_TASK_CANCEL and isinstance(payload, dict): cancel_id = payload.get("task_id", "") target_job = STATE.get_job(cancel_id) if target_job: STATE.log(f"⛔ [Worker] Cancelando tarea {cancel_id}...") if target_job.cancel_func: target_job.cancel_func() if target_job.proc and target_job.proc.poll() is None: try: STATE.log(f"💀 [Worker] Matando proceso FFmpeg PID {target_job.proc.pid} de tarea {cancel_id}...") target_job.proc.kill() except Exception: pass STATE.remove_job(cancel_id) else: STATE.log(f"⚠️ [Worker] Tarea {cancel_id} no encontrada en activeJobs para cancelar", level="WARNING") finally: ping_stop.set() ping_task.cancel() self.ws_conn = None async def _execute_task_wrapper(self, task_payload: Dict[str, Any], job: CurrentJob): """Envuelve la ejecución de la tarea en asyncio.to_thread.""" await asyncio.to_thread(self._execute_task, task_payload, job) def _execute_task(self, task: Dict[str, Any], job: CurrentJob): task_id = task.get("task_id", "") file_name = task.get("file_name", "video.mp4") STATE.log(f"📥 Nueva tarea recibida: {task_id} (Archivo: {file_name})") cancel_event = threading.Event() job.cancel_func = cancel_event.set ext = Path(file_name).suffix or ".mp4" out_ext = ".avi" if task.get("settings", {}).get("format", "").lower() == "avi" else ".mp4" input_path = os.path.join(self.work_dir, f"{task_id}_in{ext}") output_path = os.path.join(self.work_dir, f"{task_id}_out{out_ext}") wm_image_path = "" try: # 1. Descargar archivo del Manager download_url = self._resolve_url(task.get("download_url", "")) req_headers = { "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36", "X-Worker-Secret": self.secret, } job.status = "Descargando..." # El tamaño autoritativo viaja en la asignación: un proxy delante # del manager puede eliminar el Content-Length (respuesta chunked). total_size = int(task.get("input_size") or 0) asyncio.run_coroutine_threadsafe( self._send_ws_safe(TYPE_TASK_PROGRESS, { "task_id": task_id, "stage": "downloading", "percent": 0.0, "current": 0, "total": total_size, }), self.loop, ) started = time.time() received = 0 last_report = started last_report_bytes = 0 smoothed_speed = 0.0 max_attempts = 20 # Descarga por HTTP reanudable con Range: si el stream del manager # muere a mitad de la transferencia, se reabre la conexión pidiendo # los bytes que faltan en lugar de fallar la tarea entera. with open(input_path, "wb") as f: attempt = 0 while True: if cancel_event.is_set(): self._report_failed(task_id, "descarga cancelada") return if total_size and received >= total_size: break attempt += 1 headers = dict(req_headers) if received > 0: headers["Range"] = f"bytes={received}-" try: with self.session.get(download_url, headers=headers, stream=True, timeout=120) as resp: if resp.status_code == 416 and total_size and received >= total_size: break # archivo ya completo if resp.status_code == 200 and received > 0: # El servidor ignoró el Range: reescribir desde cero f.seek(0) f.truncate(0) received = 0 last_report_bytes = 0 elif resp.status_code == 206 and received > 0: f.seek(received) f.truncate(received) if total_size == 0: if resp.status_code == 206: cr = resp.headers.get("Content-Range", "") if "/" in cr: try: total_size = int(cr.rsplit("/", 1)[1]) except ValueError: pass else: total_size = int(resp.headers.get("Content-Length", "0") or 0) if resp.status_code not in (200, 206): if resp.status_code in (408, 429) or resp.status_code >= 500 or (resp.status_code == 404 and attempt <= 3): raise IOError(f"falló descarga (status {resp.status_code})") self._report_failed(task_id, f"falló descarga (status {resp.status_code})") return for chunk in resp.iter_content(chunk_size=512 * 1024): if cancel_event.is_set(): self._report_failed(task_id, "descarga cancelada") return if chunk: # Con tamaño conocido no esperamos un EOF que puede no # llegar: recortamos el último chunk y cerramos. if total_size and received + len(chunk) > total_size: chunk = chunk[: total_size - received] if chunk: f.write(chunk) received += len(chunk) now = time.time() if now - last_report >= 1 or (total_size and received >= total_size): delta_sec = max(now - last_report, 0.001) delta_bytes = received - last_report_bytes inst_speed = delta_bytes / delta_sec / (1024 * 1024) if smoothed_speed <= 0: smoothed_speed = inst_speed else: smoothed_speed = smoothed_speed * 0.7 + inst_speed * 0.3 pct = (received * 100 / total_size) if total_size else 0.0 job.percent = pct job.speed = f"{smoothed_speed:.1f} MB/s" eta = "" if total_size and received < total_size and smoothed_speed > 0: eta = f"{int((total_size - received) / (smoothed_speed * 1024 * 1024))}s" progress = { "task_id": task_id, "stage": "downloading", "percent": pct, "current": received, "total": total_size, "speed": job.speed, "eta": eta, } asyncio.run_coroutine_threadsafe( self._send_ws_safe(TYPE_TASK_PROGRESS, progress), self.loop ) last_report = now last_report_bytes = received if total_size and received >= total_size: break if total_size and received < total_size: # Conexión terminada sin excepción pero incompleta raise IOError("conexión cerrada antes de completar la descarga") f.flush() break except Exception as e: try: f.flush() except Exception: pass if attempt >= max_attempts: self._report_failed(task_id, f"falló descarga tras {max_attempts} reintentos ({received} bytes recibidos): {e}") return STATE.log( f"[{task_id}] Descarga interrumpida en {received} bytes (intento {attempt}/{max_attempts}): {e}. Reanudando...", level="WARNING", ) time.sleep(min(3 * attempt, 15)) original_size = os.path.getsize(input_path) if os.path.isfile(input_path) else 0 # Notificar al Manager que la descarga finalizó para liberar el slot de descarga asyncio.run_coroutine_threadsafe( self._send_ws_safe(TYPE_TASK_PROGRESS, {"task_id": task_id, "stage": "compressing", "percent": 0.0}), self.loop, ) # 1.5. Descargar el logo de marca de agua si la tarea incluye uno wm_url = task.get("watermark_image_url") if wm_url: candidate = os.path.join(self.work_dir, f"{task_id}_wm.png") try: with self.session.get(self._resolve_url(wm_url), headers=req_headers, timeout=30) as wm_resp: if wm_resp.status_code == 200: with open(candidate, "wb") as f: f.write(wm_resp.content) wm_image_path = candidate else: STATE.log(f"[{task_id}] No se pudo descargar el logo (status {wm_resp.status_code}); se usará marca de texto") except Exception as ex: STATE.log(f"[{task_id}] Error descargando logo ({ex}); se usará marca de texto") # 2. Obtener duración del video con ffprobe duration = get_duration(input_path) if duration <= 0: STATE.log(f"[{task_id}] Advertencia: no se pudo obtener duración exacta con ffprobe") # 3. Comprimir video con FFmpeg STATE.log(f"[{task_id}] Iniciando compresión con FFmpeg (Duración: {duration:.1f}s)...") start_t = time.time() job.status = "Comprimiendo..." def on_cmd_cb(proc: subprocess.Popen): job.proc = proc STATE.log(f"[{task_id}] Proceso FFmpeg iniciado con PID {proc.pid}") def on_progress_cb(p_data: Dict[str, Any]): job.percent = p_data.get("percent", 0.0) job.speed = p_data.get("speed", "--") job.eta = p_data.get("eta", "--") p_data["task_id"] = task_id p_data["stage"] = "compressing" asyncio.run_coroutine_threadsafe( self._send_ws_safe(TYPE_TASK_PROGRESS, p_data), self.loop, ) try: compress_with_cmd( input_path=input_path, output_path=output_path, settings=task.get("settings", {}), wm_image_path=wm_image_path, total_duration=duration, on_cmd=on_cmd_cb, on_progress=on_progress_cb, is_cancelled=cancel_event.is_set, ) except InterruptedError: STATE.log(f"[{task_id}] Compresión cancelada por solicitud del usuario/manager.") self._report_failed(task_id, "compresión cancelada") return except Exception as e: STATE.log(f"[{task_id}] Error en compresión: {e}", level="ERROR") self._report_failed(task_id, f"falló compresión: {e}") return job.done = True if not os.path.isfile(output_path) or os.path.getsize(output_path) == 0: self._report_failed(task_id, "archivo comprimido no encontrado o vacío") return compressed_size = os.path.getsize(output_path) elapsed = time.time() - start_t STATE.log(f"[{task_id}] Compresión exitosa! Original: {original_size} bytes -> Comprimido: {compressed_size} bytes ({elapsed:.1f}s)") # 4. Subir archivo comprimido al Manager en bloques (chunked resumable upload) # para evitar cortes por timeouts de reverse proxy (Caddy / Cloudflare / Nginx 60s) upload_base_url = self._resolve_url(task.get("upload_url", "")) job.status = "Subiendo..." upload_success = False chunk_size = 16 * 1024 * 1024 # 16 MB por bloque offset = 0 started = time.monotonic() asyncio.run_coroutine_threadsafe( self._send_ws_safe(TYPE_TASK_PROGRESS, { "task_id": task_id, "stage": "uploading", "percent": 0.0, "current": 0, "total": compressed_size, }), self.loop ) with open(output_path, "rb") as out_f: while offset < compressed_size: if cancel_event.is_set(): self._report_failed(task_id, "subida cancelada") return out_f.seek(offset) chunk_data = out_f.read(chunk_size) if not chunk_data: break chunk_len = len(chunk_data) end_offset = offset + chunk_len - 1 sep = "&" if "?" in upload_base_url else "?" chunk_url = f"{upload_base_url}{sep}offset={offset}&total={compressed_size}" chunk_headers = { "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36", "Content-Type": "application/octet-stream", "X-Worker-Secret": self.secret, "Content-Length": str(chunk_len), "Content-Range": f"bytes {offset}-{end_offset}/{compressed_size}", } chunk_uploaded = False for chunk_attempt in range(1, 16): if cancel_event.is_set(): self._report_failed(task_id, "subida cancelada") return try: resp = self.session.post(chunk_url, data=chunk_data, headers=chunk_headers, timeout=60) if resp.status_code in (200, 206): chunk_uploaded = True break else: STATE.log(f"[{task_id}] Manager respondió status {resp.status_code} en bloque @{offset}") time.sleep(2) except Exception as ex: if chunk_attempt < 15: STATE.log(f"[{task_id}] Reintentando bloque @{offset} ({chunk_attempt}/15): {ex}", level="WARNING") time.sleep(3) if not chunk_uploaded: STATE.log(f"[{task_id}] Falló la subida del bloque @{offset} tras múltiples reintentos", level="ERROR") break offset += chunk_len now = time.monotonic() elapsed = max(now - started, 0.001) percent = min(100.0, offset * 100.0 / max(compressed_size, 1)) speed_mb = (offset / elapsed) / (1024 * 1024) eta_sec = int((compressed_size - offset) / (offset / elapsed)) if offset < compressed_size else 0 progress = { "task_id": task_id, "stage": "uploading", "percent": percent, "current": offset, "total": compressed_size, "speed": f"{speed_mb:.1f} MB/s", "eta": f"{eta_sec}s" if eta_sec > 0 else "", } asyncio.run_coroutine_threadsafe( self._send_ws_safe(TYPE_TASK_PROGRESS, progress), self.loop, ) if offset >= compressed_size: upload_success = True if not upload_success: self._report_failed(task_id, "falló subida al manager tras múltiples reintentos") return # 5. Notificar tarea completada al Manager (reintentando si WS aún está reconectando) complete_payload = { "task_id": task_id, "original_size": original_size, "compressed_size": compressed_size, "duration_sec": round(elapsed, 2), } for _ in range(10): fut = asyncio.run_coroutine_threadsafe( self._send_ws_safe(TYPE_TASK_COMPLETE, complete_payload), self.loop, ) try: if fut.result(timeout=3): break except Exception: pass time.sleep(2) STATE.record_success(original_size, compressed_size) STATE.log(f"[{task_id}] ✅ Tarea finalizada y confirmada.") except Exception as e: # Red/FS/errores inesperados (p.ej. timeout de lectura en la # descarga): notificar el fallo para que el manager no deje la # tarea colgada indefinidamente. STATE.log(f"[{task_id}] Error inesperado en la tarea: {e}", level="ERROR") self._report_failed(task_id, f"error inesperado: {e}") finally: STATE.remove_job(task_id) for p in (input_path, output_path, wm_image_path): if p and os.path.isfile(p): try: os.remove(p) except Exception: pass def _report_failed(self, task_id: str, err_msg: str): """Notifica fallo de tarea al Manager.""" STATE.record_fail() asyncio.run_coroutine_threadsafe( self._send_ws_safe(TYPE_TASK_FAILED, {"task_id": task_id, "error": err_msg}), self.loop, )