Spaces:
Sleeping
Sleeping
| """ | |
| Custom HuggingFace Trainer Callbacks for WebSocket streaming. | |
| All async sends use asyncio.run_coroutine_threadsafe with the captured main | |
| event loop — never asyncio.new_event_loop(), which creates isolated loops | |
| that have no knowledge of the existing WebSocket connections. | |
| """ | |
| import asyncio | |
| import time | |
| from transformers import TrainerCallback, TrainingArguments, TrainerState, TrainerControl | |
| from typing import Dict, Any, Optional | |
| import logging | |
| logger = logging.getLogger(__name__) | |
| class WebSocketCallback(TrainerCallback): | |
| """ | |
| Streams training metrics to WebSocket clients in real-time. | |
| Must be constructed with the running event loop so that callbacks | |
| fired from the trainer thread can schedule coroutines on the main loop. | |
| """ | |
| def __init__( | |
| self, | |
| client_id: str, | |
| socket_manager, | |
| training_monitor=None, | |
| training_controller=None, | |
| ai_commentator=None, | |
| loop: asyncio.AbstractEventLoop | None = None, | |
| ): | |
| self.client_id = client_id | |
| self.socket_manager = socket_manager | |
| self.training_monitor = training_monitor | |
| self.training_controller = training_controller | |
| # ai_commentator kept for API compat but no longer used | |
| self._loop = loop | |
| self.start_time: float | None = None | |
| def _send_update(self, data: dict) -> None: | |
| """Schedule a WebSocket broadcast on the main event loop (fire-and-forget).""" | |
| if self._loop is None or self._loop.is_closed(): | |
| return | |
| try: | |
| asyncio.run_coroutine_threadsafe( | |
| self.socket_manager.broadcast_json(self.client_id, data), | |
| self._loop, | |
| ) | |
| except Exception as exc: | |
| logger.warning("Failed to schedule WebSocket update: %s", exc) | |
| def _schedule(self, coro) -> None: | |
| """Schedule an arbitrary coroutine on the main event loop (fire-and-forget).""" | |
| if self._loop is None or self._loop.is_closed(): | |
| return | |
| try: | |
| asyncio.run_coroutine_threadsafe(coro, self._loop) | |
| except Exception as exc: | |
| logger.warning("Failed to schedule coroutine: %s", exc) | |
| def on_train_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs): | |
| self.start_time = time.time() | |
| self._send_update({ | |
| "type": "train_begin", | |
| "total_steps": state.max_steps, | |
| "num_epochs": args.num_train_epochs, | |
| "batch_size": args.per_device_train_batch_size, | |
| }) | |
| def on_epoch_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs): | |
| self._send_update({ | |
| "type": "epoch_begin", | |
| "epoch": int(state.epoch) if state.epoch else 0, | |
| }) | |
| def on_step_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs): | |
| if self.training_controller: | |
| self.training_controller.wait_if_paused() | |
| def on_step_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs): | |
| if state.global_step % 10 == 0: | |
| elapsed = time.time() - self.start_time if self.start_time else 0 | |
| self._send_update({ | |
| "type": "step", | |
| "step": state.global_step, | |
| "total_steps": state.max_steps, | |
| "epoch": int(state.epoch) if state.epoch else 0, | |
| "elapsed_time": round(elapsed, 1), | |
| }) | |
| def on_log(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, logs: Dict[str, Any] = None, **kwargs): | |
| if logs is None: | |
| return | |
| metrics: dict[str, Any] = { | |
| "type": "metrics", | |
| "step": state.global_step, | |
| "epoch": round(state.epoch, 2) if state.epoch else 0, | |
| } | |
| for key in ("loss", "learning_rate", "eval_loss", "eval_accuracy", "grad_norm"): | |
| if key in logs: | |
| metrics[key] = round(float(logs[key]), 6) | |
| if self.start_time: | |
| metrics["elapsed_time"] = round(time.time() - self.start_time, 1) | |
| self._send_update(metrics) | |
| if self.training_monitor: | |
| self._schedule(self.training_monitor.analyze_step( | |
| step=state.global_step, | |
| train_loss=logs.get("loss"), | |
| val_loss=logs.get("eval_loss"), | |
| learning_rate=logs.get("learning_rate"), | |
| grad_norm=logs.get("grad_norm"), | |
| )) | |
| def on_epoch_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs): | |
| self._send_update({ | |
| "type": "epoch_end", | |
| "epoch": int(state.epoch) if state.epoch else 0, | |
| }) | |
| def on_train_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs): | |
| elapsed = time.time() - self.start_time if self.start_time else 0 | |
| self._send_update({ | |
| "type": "train_end", | |
| "total_time": round(elapsed, 1), | |
| "final_step": state.global_step, | |
| }) | |