trainpulse 0.2.0__py3-none-any.whl

trainpulse/__init__.py

@@ -0,0 +1,33 @@
+"""trainpulse — lightweight training health monitor."""
+
+__version__ = "0.2.0"
+
+from trainpulse._types import (
+    Alert,
+    AlertSeverity,
+    MetricSnapshot,
+    MetricType,
+    MonitorConfig,
+    TrainingReport,
+    TrainpulseError,
+)
+from trainpulse.callbacks import TrainingCallback
+from trainpulse.early_stopping import EarlyStopping, EarlyStopResult, recommend_patience
+from trainpulse.monitor import Monitor
+from trainpulse.wandb_callback import WandbCallback
+
+__all__ = [
+    "Alert",
+    "AlertSeverity",
+    "EarlyStopping",
+    "EarlyStopResult",
+    "MetricSnapshot",
+    "MetricType",
+    "Monitor",
+    "MonitorConfig",
+    "TrainingCallback",
+    "TrainingReport",
+    "TrainpulseError",
+    "WandbCallback",
+    "recommend_patience",
+]

trainpulse/_types.py

@@ -0,0 +1,105 @@
+"""Core types for trainpulse."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Callable, Dict, List, Optional
+
+
+class AlertSeverity(str, Enum):
+    INFO = "info"
+    WARNING = "warning"
+    CRITICAL = "critical"
+
+
+class MetricType(str, Enum):
+    LOSS = "loss"
+    GRADIENT_NORM = "gradient_norm"
+    LEARNING_RATE = "learning_rate"
+    STEP_TIME = "step_time"
+    MEMORY_USED = "memory_used"
+    CUSTOM = "custom"
+
+
+@dataclass
+class MetricSnapshot:
+    """A single metric recording at a given step."""
+
+    step: int
+    name: str
+    value: float
+    metric_type: MetricType = MetricType.CUSTOM
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class Alert:
+    """An alert triggered by a detector."""
+
+    step: int
+    severity: AlertSeverity
+    detector: str
+    message: str
+    metric_name: str = ""
+    metric_value: float = 0.0
+
+    def __str__(self) -> str:
+        return f"[{self.severity.value.upper()}] Step {self.step}: {self.message}"
+
+
+@dataclass
+class MonitorConfig:
+    """Configuration for the training monitor."""
+
+    # Loss spike detection
+    loss_spike_threshold: float = 5.0  # Multiplier over rolling average
+    loss_spike_window: int = 50  # Rolling window size
+
+    # Gradient monitoring
+    grad_norm_threshold: float = 100.0  # Max acceptable gradient norm
+    grad_vanish_threshold: float = 1e-7  # Min acceptable gradient norm
+
+    # NaN/Inf detection
+    check_nan: bool = True
+
+    # Learning rate monitoring
+    lr_change_threshold: float = 10.0  # Max acceptable LR ratio change per step
+
+    # Step time monitoring
+    step_time_spike_threshold: float = 3.0  # Multiplier over rolling average
+    step_time_window: int = 20
+
+    # Plateau detection
+    plateau_patience: int = 100  # Steps without improvement
+    plateau_min_delta: float = 1e-5  # Minimum change to count as improvement
+
+    # General
+    log_interval: int = 1  # Record every N steps
+    alert_callbacks: List[Callable[[Alert], None]] = field(default_factory=list)
+
+
+@dataclass
+class TrainingReport:
+    """Summary report of training health."""
+
+    total_steps: int
+    alerts: List[Alert]
+    metrics_summary: Dict[str, Dict[str, float]]  # metric_name -> {min, max, mean, last}
+    health_score: float  # 0.0 (terrible) to 1.0 (perfect)
+
+    @property
+    def n_warnings(self) -> int:
+        return sum(1 for a in self.alerts if a.severity == AlertSeverity.WARNING)
+
+    @property
+    def n_critical(self) -> int:
+        return sum(1 for a in self.alerts if a.severity == AlertSeverity.CRITICAL)
+
+    @property
+    def is_healthy(self) -> bool:
+        return self.n_critical == 0
+
+
+class TrainpulseError(Exception):
+    """Base exception."""

trainpulse/callbacks.py

@@ -0,0 +1,175 @@
+"""Framework integration callbacks."""
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+from trainpulse._types import MonitorConfig
+from trainpulse.monitor import Monitor
+
+
+class TrainingCallback:
+    """Generic callback that wraps a Monitor for any training loop.
+
+    Usage::
+
+        cb = TrainingCallback()
+        for step in range(num_steps):
+            cb.on_step_begin(step)
+            loss = train_step()
+            grad_norm = get_grad_norm()
+            cb.on_step_end(step, loss=loss, grad_norm=grad_norm, lr=optimizer.lr)
+        report = cb.report()
+    """
+
+    def __init__(self, config: Optional[MonitorConfig] = None) -> None:
+        self.monitor = Monitor(config)
+
+    def on_step_begin(self, step: int) -> None:
+        self.monitor.step_start()
+
+    def on_step_end(
+        self,
+        step: int,
+        loss: Optional[float] = None,
+        grad_norm: Optional[float] = None,
+        lr: Optional[float] = None,
+        **extra_metrics: float,
+    ) -> None:
+        self.monitor.step_end(step)
+        if loss is not None:
+            self.monitor.log("loss", step, loss)
+        if grad_norm is not None:
+            self.monitor.log("grad_norm", step, grad_norm)
+        if lr is not None:
+            self.monitor.log("learning_rate", step, lr)
+        for name, value in extra_metrics.items():
+            self.monitor.log(name, step, value)
+
+    def report(self) -> Any:
+        return self.monitor.report()
+
+
+def make_pytorch_hooks(
+    model: Any,
+    monitor: Monitor,
+) -> list:
+    """Register backward hooks on a PyTorch model to track gradient norms.
+
+    Returns a list of hook handles for cleanup.
+
+    Usage::
+
+        monitor = Monitor()
+        hooks = make_pytorch_hooks(model, monitor)
+        # ... training loop ...
+        for h in hooks:
+            h.remove()
+    """
+    try:
+        import torch  # type: ignore[import-untyped]
+    except ImportError:
+        raise ImportError("PyTorch is required: pip install trainpulse[torch]")
+
+    handles = []
+    _step_counter = {"step": 0}
+
+    def _grad_hook(module: Any, grad_input: Any, grad_output: Any) -> None:
+        total_norm = 0.0
+        for p in module.parameters():
+            if p.grad is not None:
+                total_norm += p.grad.data.norm(2).item() ** 2
+        total_norm = total_norm**0.5
+        monitor.log("grad_norm", _step_counter["step"], total_norm)
+
+    for module in model.modules():
+        # Only hook leaf modules to avoid double counting
+        if len(list(module.children())) == 0:
+            h = module.register_full_backward_hook(_grad_hook)
+            handles.append(h)
+
+    class _HookManager:
+        """Manages hooks and step counter."""
+
+        def __init__(self, handles: list, counter: dict) -> None:
+            self._handles = handles
+            self._counter = counter
+
+        def set_step(self, step: int) -> None:
+            self._counter["step"] = step
+
+        def remove(self) -> None:
+            for h in self._handles:
+                h.remove()
+
+    manager = _HookManager(handles, _step_counter)
+    return [manager]
+
+
+def make_hf_callback(
+    config: Optional[MonitorConfig] = None,
+) -> Any:
+    """Create a HuggingFace Trainer callback.
+
+    Usage::
+
+        from transformers import Trainer
+        from trainpulse.callbacks import make_hf_callback
+
+        cb = make_hf_callback()
+        trainer = Trainer(..., callbacks=[cb])
+        trainer.train()
+        report = cb.trainpulse_monitor.report()
+
+    Returns a TrainerCallback subclass instance.
+    """
+    try:
+        from transformers import TrainerCallback, TrainerControl, TrainerState, TrainingArguments  # type: ignore[import-untyped]
+    except ImportError:
+        raise ImportError(
+            "HuggingFace transformers is required: pip install transformers"
+        )
+
+    monitor = Monitor(config)
+
+    class _TrainpulseCallback(TrainerCallback):
+        def __init__(self) -> None:
+            self.trainpulse_monitor = monitor
+
+        def on_step_begin(
+            self,
+            args: TrainingArguments,
+            state: TrainerState,
+            control: TrainerControl,
+            **kwargs: Any,
+        ) -> None:
+            monitor.step_start()
+
+        def on_log(
+            self,
+            args: TrainingArguments,
+            state: TrainerState,
+            control: TrainerControl,
+            logs: Optional[dict] = None,
+            **kwargs: Any,
+        ) -> None:
+            if logs is None:
+                return
+            step = state.global_step
+            if "loss" in logs:
+                monitor.log("loss", step, logs["loss"])
+            if "learning_rate" in logs:
+                monitor.log("learning_rate", step, logs["learning_rate"])
+            if "grad_norm" in logs:
+                monitor.log("grad_norm", step, logs["grad_norm"])
+
+        def on_step_end(
+            self,
+            args: TrainingArguments,
+            state: TrainerState,
+            control: TrainerControl,
+            **kwargs: Any,
+        ) -> None:
+            monitor.step_end(state.global_step)
+
+    return _TrainpulseCallback()

trainpulse/cli.py

@@ -0,0 +1,125 @@
+"""CLI for trainpulse."""
+
+from __future__ import annotations
+
+import json
+import sys
+from pathlib import Path
+from typing import Optional
+
+
+def _build_cli():  # type: ignore[no-untyped-def]
+    """Build the CLI. Deferred import so click/rich are optional."""
+    try:
+        import click
+    except ImportError:
+        raise SystemExit(
+            "CLI dependencies required: pip install trainpulse[cli]"
+        )
+
+    @click.group()
+    @click.version_option(package_name="trainpulse")
+    def cli() -> None:
+        """trainpulse — lightweight training health monitor."""
+
+    @cli.command()
+    @click.argument("log_file", type=click.Path(exists=True))
+    @click.option("--json-out", "-o", type=click.Path(), default=None, help="Save report as JSON.")
+    @click.option("--loss-key", default="loss", help="Key for loss in log entries.")
+    @click.option("--grad-key", default="grad_norm", help="Key for gradient norm.")
+    @click.option("--lr-key", default="learning_rate", help="Key for learning rate.")
+    @click.option("--step-key", default="step", help="Key for step number.")
+    def analyze(
+        log_file: str,
+        json_out: Optional[str],
+        loss_key: str,
+        grad_key: str,
+        lr_key: str,
+        step_key: str,
+    ) -> None:
+        """Analyze a training log file (JSONL format).
+
+        Each line should be a JSON object with at least a step and loss field.
+        """
+        from trainpulse.monitor import Monitor
+        from trainpulse.report import format_report_rich, format_report_text, save_json
+
+        monitor = Monitor()
+        path = Path(log_file)
+
+        n_lines = 0
+        for line in path.read_text().splitlines():
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                entry = json.loads(line)
+            except json.JSONDecodeError:
+                continue
+
+            step = entry.get(step_key, n_lines)
+
+            if loss_key in entry:
+                monitor.log("loss", step, float(entry[loss_key]))
+            if grad_key in entry:
+                monitor.log("grad_norm", step, float(entry[grad_key]))
+            if lr_key in entry:
+                monitor.log("learning_rate", step, float(entry[lr_key]))
+
+            n_lines += 1
+
+        if n_lines == 0:
+            click.echo("No log entries found.", err=True)
+            sys.exit(1)
+
+        report = monitor.report()
+
+        try:
+            output = format_report_rich(report)
+        except Exception:
+            output = format_report_text(report)
+        click.echo(output)
+
+        if json_out:
+            save_json(report, json_out)
+            click.echo(f"Report saved to {json_out}")
+
+    @cli.command()
+    @click.argument("report_file", type=click.Path(exists=True))
+    def show(report_file: str) -> None:
+        """Display a previously saved JSON report."""
+        from trainpulse._types import Alert, AlertSeverity, TrainingReport
+        from trainpulse.report import format_report_rich, format_report_text, load_json
+
+        data = load_json(report_file)
+        alerts = [
+            Alert(
+                step=a["step"],
+                severity=AlertSeverity(a["severity"]),
+                detector=a["detector"],
+                message=a["message"],
+                metric_name=a.get("metric_name", ""),
+                metric_value=a.get("metric_value", 0.0),
+            )
+            for a in data.get("alerts", [])
+        ]
+        report = TrainingReport(
+            total_steps=data["total_steps"],
+            alerts=alerts,
+            metrics_summary=data.get("metrics_summary", {}),
+            health_score=data.get("health_score", 1.0),
+        )
+
+        try:
+            output = format_report_rich(report)
+        except Exception:
+            output = format_report_text(report)
+        click.echo(output)
+
+    return cli
+
+
+cli = _build_cli()
+
+if __name__ == "__main__":
+    cli()

trainpulse/detectors.py

@@ -0,0 +1,216 @@
+"""Anomaly detectors for training metrics."""
+
+from __future__ import annotations
+
+import math
+from typing import List, Optional, Sequence
+
+from trainpulse._types import Alert, AlertSeverity, MetricType
+
+
+class RollingWindow:
+    """Fixed-size rolling window of float values."""
+
+    def __init__(self, size: int) -> None:
+        self._size = max(size, 1)
+        self._values: list[float] = []
+
+    def add(self, value: float) -> None:
+        self._values.append(value)
+        if len(self._values) > self._size:
+            self._values.pop(0)
+
+    @property
+    def values(self) -> List[float]:
+        return list(self._values)
+
+    @property
+    def mean(self) -> float:
+        if not self._values:
+            return 0.0
+        return sum(self._values) / len(self._values)
+
+    @property
+    def std(self) -> float:
+        if len(self._values) < 2:
+            return 0.0
+        m = self.mean
+        return math.sqrt(sum((v - m) ** 2 for v in self._values) / len(self._values))
+
+    @property
+    def is_full(self) -> bool:
+        return len(self._values) >= self._size
+
+    def __len__(self) -> int:
+        return len(self._values)
+
+
+class NaNDetector:
+    """Detect NaN or Inf values in metrics."""
+
+    def check(self, step: int, name: str, value: float) -> Optional[Alert]:
+        if math.isnan(value):
+            return Alert(
+                step=step,
+                severity=AlertSeverity.CRITICAL,
+                detector="nan_detector",
+                message=f"{name} is NaN",
+                metric_name=name,
+                metric_value=value,
+            )
+        if math.isinf(value):
+            return Alert(
+                step=step,
+                severity=AlertSeverity.CRITICAL,
+                detector="nan_detector",
+                message=f"{name} is Inf",
+                metric_name=name,
+                metric_value=value,
+            )
+        return None
+
+
+class LossSpikeDetector:
+    """Detect sudden spikes in loss values."""
+
+    def __init__(self, threshold: float = 5.0, window_size: int = 50) -> None:
+        self._threshold = threshold
+        self._window = RollingWindow(window_size)
+
+    def check(self, step: int, value: float) -> Optional[Alert]:
+        if self._window.is_full:
+            avg = self._window.mean
+            if avg > 0 and value > avg * self._threshold:
+                alert = Alert(
+                    step=step,
+                    severity=AlertSeverity.WARNING,
+                    detector="loss_spike",
+                    message=f"Loss spike: {value:.4f} ({value/avg:.1f}x rolling avg {avg:.4f})",
+                    metric_name="loss",
+                    metric_value=value,
+                )
+                self._window.add(value)
+                return alert
+
+        self._window.add(value)
+        return None
+
+
+class GradientDetector:
+    """Detect gradient explosion and vanishing."""
+
+    def __init__(
+        self,
+        explosion_threshold: float = 100.0,
+        vanish_threshold: float = 1e-7,
+    ) -> None:
+        self._explosion = explosion_threshold
+        self._vanish = vanish_threshold
+
+    def check(self, step: int, grad_norm: float) -> Optional[Alert]:
+        if grad_norm > self._explosion:
+            return Alert(
+                step=step,
+                severity=AlertSeverity.CRITICAL,
+                detector="gradient",
+                message=f"Gradient explosion: norm={grad_norm:.4f} (threshold={self._explosion})",
+                metric_name="gradient_norm",
+                metric_value=grad_norm,
+            )
+        if 0 < grad_norm < self._vanish:
+            return Alert(
+                step=step,
+                severity=AlertSeverity.WARNING,
+                detector="gradient",
+                message=f"Vanishing gradient: norm={grad_norm:.2e} (threshold={self._vanish:.2e})",
+                metric_name="gradient_norm",
+                metric_value=grad_norm,
+            )
+        return None
+
+
+class LRDetector:
+    """Detect suspicious learning rate changes."""
+
+    def __init__(self, change_threshold: float = 10.0) -> None:
+        self._threshold = change_threshold
+        self._prev_lr: Optional[float] = None
+
+    def check(self, step: int, lr: float) -> Optional[Alert]:
+        if self._prev_lr is not None and self._prev_lr > 0 and lr > 0:
+            ratio = max(lr / self._prev_lr, self._prev_lr / lr)
+            if ratio > self._threshold:
+                alert = Alert(
+                    step=step,
+                    severity=AlertSeverity.WARNING,
+                    detector="learning_rate",
+                    message=f"LR changed {ratio:.1f}x in one step ({self._prev_lr:.2e} → {lr:.2e})",
+                    metric_name="learning_rate",
+                    metric_value=lr,
+                )
+                self._prev_lr = lr
+                return alert
+        self._prev_lr = lr
+        return None
+
+
+class PlateauDetector:
+    """Detect loss plateaus (no improvement for N steps)."""
+
+    def __init__(self, patience: int = 100, min_delta: float = 1e-5) -> None:
+        self._patience = patience
+        self._min_delta = min_delta
+        self._best_loss: Optional[float] = None
+        self._steps_without_improvement = 0
+        self._alerted = False
+
+    def check(self, step: int, loss: float) -> Optional[Alert]:
+        if self._best_loss is None:
+            self._best_loss = loss
+            return None
+
+        if loss < self._best_loss - self._min_delta:
+            self._best_loss = loss
+            self._steps_without_improvement = 0
+            self._alerted = False
+            return None
+
+        self._steps_without_improvement += 1
+
+        if self._steps_without_improvement >= self._patience and not self._alerted:
+            self._alerted = True
+            return Alert(
+                step=step,
+                severity=AlertSeverity.WARNING,
+                detector="plateau",
+                message=f"Loss plateau: no improvement for {self._steps_without_improvement} steps (best={self._best_loss:.6f})",
+                metric_name="loss",
+                metric_value=loss,
+            )
+        return None
+
+
+class StepTimeDetector:
+    """Detect unusually slow training steps."""
+
+    def __init__(self, threshold: float = 3.0, window_size: int = 20) -> None:
+        self._threshold = threshold
+        self._window = RollingWindow(window_size)
+
+    def check(self, step: int, step_time: float) -> Optional[Alert]:
+        if self._window.is_full:
+            avg = self._window.mean
+            if avg > 0 and step_time > avg * self._threshold:
+                alert = Alert(
+                    step=step,
+                    severity=AlertSeverity.WARNING,
+                    detector="step_time",
+                    message=f"Slow step: {step_time:.2f}s ({step_time/avg:.1f}x avg {avg:.2f}s)",
+                    metric_name="step_time",
+                    metric_value=step_time,
+                )
+                self._window.add(step_time)
+                return alert
+
+        self._window.add(step_time)
+        return None