DeepFabric 4.4.0__py3-none-any.whl

deepfabric/training/callback.py

@@ -0,0 +1,363 @@
+"""DeepFabric TrainerCallback for automatic metrics logging."""
+
+from __future__ import annotations
+
+import logging
+import os
+import uuid
+
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any
+
+from .api_key_prompt import get_api_key
+from .metrics_sender import MetricsSender
+
+if TYPE_CHECKING:
+    from transformers import TrainerControl, TrainerState
+    from transformers.training_args import TrainingArguments
+
+logger = logging.getLogger(__name__)
+
+
+class DeepFabricCallback:
+    """Callback that sends training metrics to DeepFabric SaaS.
+
+    This callback integrates with HuggingFace Trainer and TRL trainers to
+    automatically log training metrics (loss, learning rate, epoch, global step,
+    throughput, TRL-specific metrics, and custom metrics) to the DeepFabric
+    backend.
+
+    The callback is designed to be non-blocking and gracefully handles failures
+    without impacting training.
+
+    Example:
+        from deepfabric.training import DeepFabricCallback
+
+        trainer = Trainer(
+            model=model,
+            args=training_args,
+            train_dataset=train_dataset,
+        )
+        trainer.add_callback(DeepFabricCallback(trainer))
+        trainer.train()
+
+    Environment Variables:
+        DEEPFABRIC_API_KEY: API key (alternative to constructor arg)
+        DEEPFABRIC_API_URL: Backend URL (default: https://api.deepfabric.ai)
+    """
+
+    def __init__(
+        self,
+        trainer: Any | None = None,
+        api_key: str | None = None,
+        endpoint: str | None = None,
+        enabled: bool = True,
+    ):
+        """Initialize the DeepFabric callback.
+
+        Args:
+            trainer: Optional Trainer instance to extract model info from
+            api_key: DeepFabric API key (falls back to DEEPFABRIC_API_KEY env var,
+                     then prompts in interactive environments)
+            endpoint: API endpoint URL (falls back to DEEPFABRIC_API_URL env var)
+            enabled: Whether logging is enabled (default: True)
+        """
+        # Get API key from arg, env, or prompt
+        self.api_key = api_key or get_api_key()
+        self.endpoint = endpoint or os.getenv("DEEPFABRIC_API_URL", "https://api.deepfabric.ai")
+        self.run_id = str(uuid.uuid4())
+        self.enabled = enabled and self.api_key is not None
+
+        # Store trainer reference for model extraction
+        self._trainer = trainer
+
+        # Initialize sender (handles None api_key gracefully)
+        self.sender = MetricsSender(
+            endpoint=self.endpoint,
+            api_key=self.api_key if self.enabled else None,
+        )
+
+        self._run_started = False
+        self._model_name: str | None = None
+        self._training_args_logged = False
+
+        if self.enabled:
+            logger.debug(f"DeepFabric callback initialized (run_id={self.run_id})")
+        else:
+            logger.debug("DeepFabric callback disabled (no API key)")
+
+    def on_train_begin(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,  # noqa: ARG002
+        **kwargs: Any,
+    ) -> None:
+        """Called at the beginning of training.
+
+        Sends run start event with training configuration.
+        """
+        if not self.enabled or self._run_started:
+            return
+
+        self._run_started = True
+
+        # Extract model name from various sources
+        model = kwargs.get("model")
+        if model is None and self._trainer is not None:
+            model = getattr(self._trainer, "model", None)
+        self._model_name = self._extract_model_name(args, model)
+
+        # Build training args dict (safe extraction)
+        training_config = self._extract_training_args(args)
+
+        self.sender.send_run_start(
+            {
+                "run_id": self.run_id,
+                "model_name": self._model_name,
+                "training_config": training_config,
+                "state": {
+                    "max_steps": state.max_steps,
+                    "num_train_epochs": state.num_train_epochs,
+                    "is_world_process_zero": getattr(state, "is_world_process_zero", True),
+                },
+            }
+        )
+
+    def on_log(
+        self,
+        args: TrainingArguments,  # noqa: ARG002
+        state: TrainerState,
+        control: TrainerControl,  # noqa: ARG002
+        logs: dict[str, float] | None = None,
+        **kwargs: Any,  # noqa: ARG002
+    ) -> None:
+        """Called when metrics are logged.
+
+        Sends all logged metrics to DeepFabric (loss, learning_rate, epoch,
+        global_step, throughput, TRL metrics, custom metrics, etc.).
+        """
+        if not self.enabled or logs is None:
+            return
+
+        # Filter out None values and non-numeric values
+        filtered_logs = {}
+        for key, value in logs.items():
+            if value is not None:
+                if isinstance(value, int | float):
+                    filtered_logs[key] = value
+                elif isinstance(value, str):
+                    # Keep string values for metadata
+                    filtered_logs[key] = value
+
+        if not filtered_logs:
+            return
+
+        payload = {
+            "run_id": self.run_id,
+            "global_step": state.global_step,
+            "epoch": state.epoch,
+            "timestamp": datetime.now(timezone.utc).isoformat(),
+            "type": "log",
+            "metrics": filtered_logs,
+        }
+
+        self.sender.send_metrics(payload)
+
+    def on_evaluate(
+        self,
+        args: TrainingArguments,  # noqa: ARG002
+        state: TrainerState,
+        control: TrainerControl,  # noqa: ARG002
+        metrics: dict[str, float] | None = None,
+        **kwargs: Any,  # noqa: ARG002
+    ) -> None:
+        """Called after evaluation.
+
+        Sends evaluation metrics to DeepFabric.
+        """
+        if not self.enabled or metrics is None:
+            return
+
+        payload = {
+            "run_id": self.run_id,
+            "global_step": state.global_step,
+            "epoch": state.epoch,
+            "timestamp": datetime.now(timezone.utc).isoformat(),
+            "type": "eval",
+            "metrics": metrics,
+        }
+
+        self.sender.send_metrics(payload)
+
+    def on_train_end(
+        self,
+        args: TrainingArguments,  # noqa: ARG002
+        state: TrainerState,
+        control: TrainerControl,  # noqa: ARG002
+        **kwargs: Any,  # noqa: ARG002
+    ) -> None:
+        """Called at the end of training.
+
+        Sends run end event and flushes pending metrics.
+        """
+        if not self.enabled or not self._run_started:
+            return
+
+        self.sender.send_run_end(
+            {
+                "run_id": self.run_id,
+                "final_step": state.global_step,
+                "final_epoch": state.epoch,
+                "total_flos": getattr(state, "total_flos", None),
+                "best_metric": getattr(state, "best_metric", None),
+                "best_model_checkpoint": getattr(state, "best_model_checkpoint", None),
+            }
+        )
+
+        # Flush remaining metrics
+        self.sender.flush(timeout=30.0)
+
+        logger.debug(f"DeepFabric run completed: {self.sender.stats}")
+
+    def on_save(
+        self,
+        args: TrainingArguments,  # noqa: ARG002
+        state: TrainerState,
+        control: TrainerControl,  # noqa: ARG002
+        **kwargs: Any,  # noqa: ARG002
+    ) -> None:
+        """Called when a checkpoint is saved.
+
+        Optionally logs checkpoint events.
+        """
+        if not self.enabled:
+            return
+
+        # Log checkpoint event
+        self.sender.send_metrics(
+            {
+                "run_id": self.run_id,
+                "global_step": state.global_step,
+                "epoch": state.epoch,
+                "timestamp": datetime.now(timezone.utc).isoformat(),
+                "type": "checkpoint",
+                "metrics": {"checkpoint_step": state.global_step},
+            }
+        )
+
+    def _extract_model_name(self, args: TrainingArguments, model: Any | None) -> str | None:
+        """Extract model name from various sources.
+
+        Args:
+            args: Training arguments
+            model: Model instance (may be None)
+
+        Returns:
+            Model name or None
+        """
+        # Try args first
+        if hasattr(args, "model_name_or_path"):
+            return args.model_name_or_path
+
+        # Try model config
+        if model is not None:
+            if hasattr(model, "config") and hasattr(model.config, "name_or_path"):
+                return model.config.name_or_path
+            if hasattr(model, "name_or_path"):
+                return model.name_or_path
+
+        # Try output_dir as fallback
+        if hasattr(args, "output_dir"):
+            return os.path.basename(args.output_dir)
+
+        return None
+
+    def _extract_training_args(self, args: TrainingArguments) -> dict[str, Any]:
+        """Extract training arguments for logging.
+
+        Args:
+            args: Training arguments
+
+        Returns:
+            Dictionary of training configuration
+        """
+        config = {}
+
+        # Core training args
+        safe_attrs = [
+            "num_train_epochs",
+            "max_steps",
+            "per_device_train_batch_size",
+            "per_device_eval_batch_size",
+            "gradient_accumulation_steps",
+            "learning_rate",
+            "weight_decay",
+            "adam_beta1",
+            "adam_beta2",
+            "adam_epsilon",
+            "max_grad_norm",
+            "warmup_steps",
+            "warmup_ratio",
+            "lr_scheduler_type",
+            "logging_steps",
+            "eval_steps",
+            "save_steps",
+            "seed",
+            "fp16",
+            "bf16",
+            "gradient_checkpointing",
+            "deepspeed",
+            "local_rank",
+            "dataloader_num_workers",
+        ]
+
+        for attr in safe_attrs:
+            if hasattr(args, attr):
+                value = getattr(args, attr)
+                # Convert enums to strings
+                if hasattr(value, "value"):
+                    value = value.value
+                config[attr] = value
+
+        return config
+
+
+# Make it compatible with transformers TrainerCallback protocol
+# by ensuring it has all required methods (even as no-ops)
+def _ensure_trainer_callback_compatibility():
+    """Ensure DeepFabricCallback has all TrainerCallback methods."""
+    # These methods are optional but good to have for completeness
+    # Include all methods that transformers might call on callbacks
+    optional_methods = [
+        "on_step_begin",
+        "on_step_end",
+        "on_substep_end",
+        "on_epoch_begin",
+        "on_epoch_end",
+        "on_prediction_step",
+        "on_init_end",
+        # Newer transformers versions
+        "on_pre_optimizer_step",
+        "on_optimizer_step",
+        "on_post_optimizer_step",
+        "on_pre_scheduler_step",
+        "on_scheduler_step",
+        "on_post_scheduler_step",
+    ]
+
+    def _make_noop(name):
+        """Create a no-op method that returns control unchanged."""
+
+        def noop(self, args, state, control, **kwargs):  # noqa: ARG001
+            return control
+
+        noop.__name__ = name
+        return noop
+
+    for method in optional_methods:
+        if not hasattr(DeepFabricCallback, method):
+            setattr(DeepFabricCallback, method, _make_noop(method))
+
+
+_ensure_trainer_callback_compatibility()

deepfabric/training/metrics_sender.py

@@ -0,0 +1,301 @@
+"""Non-blocking async metrics sender for training metrics."""
+
+from __future__ import annotations
+
+import atexit
+import logging
+import queue
+import threading
+import time
+
+from datetime import datetime, timezone
+from typing import Any
+
+import requests
+
+logger = logging.getLogger(__name__)
+
+
+class MetricsSender:
+    """Non-blocking metrics sender with background thread.
+
+    Queues metrics and sends them in batches via a background thread to avoid
+    blocking training. Gracefully handles network errors and queue overflow.
+
+    Example:
+        sender = MetricsSender(
+            endpoint="https://api.deepfabric.ai",
+            api_key="your-api-key",
+        )
+        sender.send_metrics({"loss": 2.5, "step": 100})
+        sender.flush()  # Ensure all metrics are sent
+    """
+
+    def __init__(
+        self,
+        endpoint: str,
+        api_key: str | None,
+        batch_size: int = 10,
+        flush_interval: float = 5.0,
+        max_queue_size: int = 1000,
+        timeout: float = 10.0,
+    ):
+        """Initialize the metrics sender.
+
+        Args:
+            endpoint: Base URL for the DeepFabric API
+            api_key: API key for authentication (None disables sending)
+            batch_size: Number of metrics to batch before sending
+            flush_interval: Seconds between automatic flushes
+            max_queue_size: Maximum queue size (overflow drops metrics)
+            timeout: HTTP request timeout in seconds
+        """
+        self.endpoint = endpoint.rstrip("/")
+        self.api_key = api_key
+        self.batch_size = batch_size
+        self.flush_interval = flush_interval
+        self.timeout = timeout
+
+        self._queue: queue.Queue[dict[str, Any]] = queue.Queue(maxsize=max_queue_size)
+        self._stop_event = threading.Event()
+        self._enabled = api_key is not None
+
+        # Start background sender thread
+        if self._enabled:
+            self._thread = threading.Thread(
+                target=self._sender_loop,
+                daemon=True,
+                name="deepfabric-metrics-sender",
+            )
+            self._thread.start()
+            atexit.register(self.shutdown)
+        else:
+            self._thread = None
+
+        self._send_errors = 0
+        self._metrics_sent = 0
+        self._metrics_dropped = 0
+
+    @property
+    def enabled(self) -> bool:
+        """Whether the sender is enabled (has API key)."""
+        return self._enabled
+
+    @property
+    def stats(self) -> dict[str, int]:
+        """Get sender statistics."""
+        return {
+            "metrics_sent": self._metrics_sent,
+            "metrics_dropped": self._metrics_dropped,
+            "send_errors": self._send_errors,
+            "queue_size": self._queue.qsize(),
+        }
+
+    def send_metrics(self, metrics: dict[str, Any]) -> bool:
+        """Queue metrics for async sending (non-blocking).
+
+        Args:
+            metrics: Dictionary of metric names to values
+
+        Returns:
+            True if queued successfully, False if dropped
+        """
+        if not self._enabled:
+            return False
+
+        # Add timestamp if not present
+        if "timestamp" not in metrics:
+            metrics["timestamp"] = datetime.now(timezone.utc).isoformat()
+
+        try:
+            self._queue.put_nowait({"type": "metrics", "data": metrics})
+        except queue.Full:
+            self._metrics_dropped += 1
+            logger.debug("Metrics queue full, dropping metrics")
+            return False
+        else:
+            return True
+
+    def send_run_start(self, metadata: dict[str, Any]) -> bool:
+        """Send run start event.
+
+        Args:
+            metadata: Run metadata (model_name, training_args, etc.)
+
+        Returns:
+            True if queued successfully
+        """
+        return self._send_event("run_start", metadata)
+
+    def send_run_end(self, metadata: dict[str, Any]) -> bool:
+        """Send run end event.
+
+        Args:
+            metadata: Run end metadata (final_step, final_epoch, etc.)
+
+        Returns:
+            True if queued successfully
+        """
+        return self._send_event("run_end", metadata)
+
+    def _send_event(self, event_type: str, data: dict[str, Any]) -> bool:
+        """Queue an event for sending.
+
+        Args:
+            event_type: Type of event (run_start, run_end, etc.)
+            data: Event data
+
+        Returns:
+            True if queued successfully
+        """
+        if not self._enabled:
+            return False
+
+        if "timestamp" not in data:
+            data["timestamp"] = datetime.now(timezone.utc).isoformat()
+
+        try:
+            self._queue.put_nowait({"type": event_type, "data": data})
+        except queue.Full:
+            logger.debug(f"Queue full, dropping {event_type} event")
+            return False
+        else:
+            return True
+
+    def _sender_loop(self) -> None:
+        """Background thread that batches and sends metrics."""
+        batch: list[dict[str, Any]] = []
+        last_flush = time.monotonic()
+
+        while not self._stop_event.is_set():
+            try:
+                # Wait for item with timeout
+                item = self._queue.get(timeout=min(1.0, self.flush_interval))
+                batch.append(item)
+
+                # Check if we should flush
+                should_flush = (
+                    len(batch) >= self.batch_size
+                    or (time.monotonic() - last_flush) >= self.flush_interval
+                )
+
+                if should_flush:
+                    self._flush_batch(batch)
+                    batch = []
+                    last_flush = time.monotonic()
+
+            except queue.Empty:
+                # Timeout - flush if we have pending items
+                if batch and (time.monotonic() - last_flush) >= self.flush_interval:
+                    self._flush_batch(batch)
+                    batch = []
+                    last_flush = time.monotonic()
+
+        # On shutdown, drain the queue and flush everything
+        while not self._queue.empty():
+            try:
+                batch.append(self._queue.get_nowait())
+            except queue.Empty:
+                break
+        if batch:
+            self._flush_batch(batch)
+
+    def _flush_batch(self, batch: list[dict[str, Any]]) -> None:
+        """Send batch of metrics to API.
+
+        Args:
+            batch: List of queued items to send
+        """
+        if not batch or not self._enabled:
+            return
+
+        # Separate events and metrics
+        events = [item for item in batch if item["type"] != "metrics"]
+        metrics = [item["data"] for item in batch if item["type"] == "metrics"]
+
+        # Send events first (run_start, run_end)
+        for event in events:
+            self._send_to_api(
+                endpoint=f"{self.endpoint}/v1/training/runs",
+                payload={"event_type": event["type"], **event["data"]},
+            )
+
+        # Send metrics batch
+        if metrics:
+            self._send_to_api(
+                endpoint=f"{self.endpoint}/v1/training/metrics",
+                payload={"metrics": metrics},
+            )
+            self._metrics_sent += len(metrics)
+
+    def _send_to_api(self, endpoint: str, payload: dict[str, Any]) -> bool:
+        """Send payload to API endpoint.
+
+        Args:
+            endpoint: Full API endpoint URL
+            payload: JSON payload to send
+
+        Returns:
+            True if sent successfully
+        """
+        try:
+            response = requests.post(
+                endpoint,
+                json=payload,
+                headers={
+                    "Authorization": f"Bearer {self.api_key}",
+                    "Content-Type": "application/json",
+                    "User-Agent": "deepfabric-training/1.0",
+                },
+                timeout=self.timeout,
+            )
+
+            if not response.ok:
+                self._send_errors += 1
+                logger.debug(f"API request failed: {response.status_code} {response.text[:100]}")
+                return False
+
+        except requests.exceptions.Timeout:
+            self._send_errors += 1
+            logger.debug("API request timed out")
+            return False
+
+        except requests.exceptions.ConnectionError:
+            self._send_errors += 1
+            logger.debug("API connection error")
+            return False
+
+        except requests.exceptions.RequestException as e:
+            self._send_errors += 1
+            logger.debug(f"API request error: {e}")
+            return False
+
+        else:
+            return True
+
+    def flush(self, timeout: float = 30.0) -> None:
+        """Flush all pending metrics (blocking).
+
+        Args:
+            timeout: Maximum time to wait for flush
+        """
+        if not self._enabled:
+            return
+
+        start = time.monotonic()
+        while not self._queue.empty() and (time.monotonic() - start) < timeout:
+            time.sleep(0.1)
+
+    def shutdown(self) -> None:
+        """Stop the sender thread and flush remaining metrics."""
+        if not self._enabled or self._thread is None:
+            return
+
+        self._stop_event.set()
+        self._thread.join(timeout=5.0)
+
+        # Log final stats
+        logger.debug(
+            f"MetricsSender shutdown: sent={self._metrics_sent}, "
+            f"dropped={self._metrics_dropped}, errors={self._send_errors}"
+        )