Module neuralops.experiment_logging

Experiment logging – structured logging, metrics tracking, and Trainer callbacks.

Sub-modules

neuralops.experiment_logging.callbacks

HuggingFace Trainer callback for automatic metric logging.

neuralops.experiment_logging.config

Logging configuration dataclass.

neuralops.experiment_logging.integration

Automatic integration between logging, ResultsSaver, and the experiment runner. Call auto_setup()(run_config) once at the start of training.

neuralops.experiment_logging.logger

Structured and thread-safe logger setup with JSON and human-readable formats. Implements QueueHandler / QueueListener for non-blocking logging across …

neuralops.experiment_logging.tracker

Simple in‑memory metrics tracker with JSON serialisation.

Functions

def auto_setup(run_config: dict, run_id: str, base_dir: pathlib.Path) ‑> dict

Expand source code

def auto_setup(
    run_config: dict,
    run_id: str,
    base_dir: Path,
) -> dict:
    """
    Create logger and tracker from run config, returning them as a dict.
    Also writes config snapshot to the run directory.

    Args:
        run_config: dictionary containing keys such as `log_level`, `log_format`,
                    `save_logs`, `exp_name`, etc.
        run_id: unique identifier for this run.
        base_dir: root directory where run artifacts are stored.

    Returns:
        A dictionary with keys `logger`, `tracker`, and `config`.
    """
    log_cfg = LoggingConfig(
        log_level=run_config.get("log_level", "INFO"),
        log_format=run_config.get("log_format", "structured"),
        log_to_console=True,
        log_to_file=run_config.get("save_logs", True),
        log_dir=str(base_dir / run_id / "logs"),
        experiment_name=run_config.get("exp_name", "training"),
        run_id=run_id,
        use_timestamp=True,
        metrics_enabled=True,
    )
    logger = setup_logging(log_cfg)

    metrics_path = base_dir / run_id / "logs" / "metrics.jsonl"
    tracker = MetricsTracker(str(metrics_path))

    # Save config snapshot for reproducibility
    import json

    config_dir = base_dir / run_id / "artifacts"
    config_dir.mkdir(parents=True, exist_ok=True)
    with open(config_dir / "config.json", "w") as f:
        json.dump(run_config, f, indent=2, default=str)

    logger.info("Auto‑setup complete. Run ID: %s", run_id)
    return {"logger": logger, "tracker": tracker, "config": log_cfg}

Create logger and tracker from run config, returning them as a dict. Also writes config snapshot to the run directory.

Args

-----=

run_config

dictionary containing keys such as log_level, log_format, save_logs, exp_name, etc.

run_id

unique identifier for this run.

base_dir

root directory where run artifacts are stored.

Returns -----= A dictionary with keys neuralops.experiment_logging.logger, neuralops.experiment_logging.tracker, and neuralops.experiment_logging.config.

def get_logger(name: str = 'neuralops') ‑> logging.Logger

Expand source code

def get_logger(name: str = "neuralops") -> logging.Logger:
    """Return the logger."""
    return logging.getLogger(name)

Return the logger.

def get_rank(configured_rank: int | None = None) ‑> int

Expand source code

def get_rank(configured_rank: int | None = None) -> int:
    """
    Determine the current process rank in single-node or DDP multi-GPU environments.

    Args:
        configured_rank: Explicit rank if specified in config.

    Returns:
        int: Process local/global rank (0 for single-node / master).
    """
    if configured_rank is not None:
        return int(configured_rank)

    for env_key in (
        "LOCAL_RANK",
        "RANK",
        "SLURM_LOCALID",
        "OMPI_COMM_WORLD_LOCAL_RANK",
    ):
        val = os.environ.get(env_key)
        if val is not None:
            try:
                return int(val)
            except ValueError:
                pass

    try:
        import torch.distributed as dist

        if dist.is_available() and dist.is_initialized():
            return dist.get_rank()
    except (
        AttributeError,
        KeyError,
        RuntimeError,
        TypeError,
        ValueError,
        OSError,
        ImportError,
    ):
        pass

    return 0

Determine the current process rank in single-node or DDP multi-GPU environments.

Args

-----=

configured_rank

Explicit rank if specified in config.

Returns

-----=

int

Process local/global rank (0 for single-node / master).

def setup_logging(config: LoggingConfig | None = None) ‑> logging.Logger

Expand source code

def setup_logging(config: LoggingConfig | None = None) -> logging.Logger:
    """
    Create and configure a thread-safe logger with isolated rank logging.

    Args:
        config: LoggingConfig dataclass specifying log parameters.

    Returns:
        logging.Logger: Configured logger instance.
    """
    cfg = config or LoggingConfig()
    rank = get_rank(cfg.rank)

    root = logging.getLogger()
    root.handlers.clear()
    root.setLevel(getattr(logging, cfg.log_level.upper(), logging.INFO))

    os.makedirs(cfg.log_dir, exist_ok=True)

    if cfg.log_format == "structured":
        fmt = StructuredFormatter(rank=rank)
    else:
        fmt = logging.Formatter(
            f"%(asctime)s [Rank {rank}] [%(levelname)s] %(name)s: %(message)s",
            datefmt="%Y-%m-%d %H:%M:%S",
        )

    destination_handlers: list[logging.Handler] = []

    if cfg.log_to_console:
        # Stream logs to stdout (typically rank 0 only or all ranks if formatted)
        ch = logging.StreamHandler(sys.stdout)
        ch.setFormatter(fmt)
        destination_handlers.append(ch)

    if cfg.log_to_file:
        # FIX: Ensure each DDP rank writes to an isolated log file (rank_{LOCAL_RANK}.log)
        if cfg.rank_specific_logs:
            fname = f"rank_{rank}.log"
        else:
            fname = f"{cfg.experiment_name}_{cfg.run_id or ''}_rank_{rank}.log".strip("_")

        log_file_path = os.path.join(cfg.log_dir, fname)
        fh = RotatingFileHandler(
            log_file_path,
            maxBytes=10 * 1024 * 1024,
            backupCount=5,
            encoding="utf-8",
        )
        fh.setFormatter(fmt)
        destination_handlers.append(fh)

    if cfg.use_queue_handler:
        # FIX: Implement Python logging with QueueHandler and QueueListener for thread safety and non-blocking I/O
        log_queue: queue.Queue = queue.Queue(-1)
        queue_handler = QueueHandler(log_queue)
        root.addHandler(queue_handler)

        listener = QueueListener(log_queue, *destination_handlers, respect_handler_level=True)
        listener.start()
        _ACTIVE_LISTENERS.append(listener)
    else:
        for handler in destination_handlers:
            root.addHandler(handler)

    logger = logging.getLogger("neuralops")
    logger.info(
        "Logging initialised. Experiment: %s, Run: %s, Rank: %d, ThreadSafeQueue: %s",
        cfg.experiment_name,
        cfg.run_id,
        rank,
        cfg.use_queue_handler,
        extra={"rank": rank},
    )
    return logger

Create and configure a thread-safe logger with isolated rank logging.

Args

-----=

config

LoggingConfig dataclass specifying log parameters.

Returns

-----=

logging.Logger

Configured logger instance.

def shutdown_logging() ‑> None

Expand source code

def shutdown_logging() -> None:
    """
    Stop all background QueueListener workers and flush pending log events.
    """
    global _ACTIVE_LISTENERS
    for listener in _ACTIVE_LISTENERS:
        try:
            listener.stop()
        except (
            AttributeError,
            KeyError,
            RuntimeError,
            TypeError,
            ValueError,
            OSError,
            ImportError,
        ):
            pass
    _ACTIVE_LISTENERS.clear()

Stop all background QueueListener workers and flush pending log events.

Classes

class LoggingCallback (tracker: MetricsTracker,
logger: logging.Logger)

Expand source code

class LoggingCallback(TrainerCallback):
    """
    Automatically captures training/evaluation logs and stores them in a MetricsTracker.

    Usage:
        tracker = MetricsTracker("metrics.json")
        trainer.add_callback(LoggingCallback(tracker, logger))
    """

    def __init__(self, tracker: MetricsTracker, logger: logging.Logger) -> None:
        self.tracker = tracker
        self.logger = logger

    def on_log(self, args, state, control, logs=None, **kwargs) -> None:
        """On log."""
        if logs:
            step = state.global_step
            self.tracker.log(logs, step=step)
            # Log key metrics to the console/file logger
            loss = logs.get("loss", None)
            eval_loss = logs.get("eval_loss", None)
            lr = logs.get("learning_rate", None)
            msg = f"Step {step}: "
            if loss is not None:
                msg += f"loss={loss:.4f} "
            if eval_loss is not None:
                msg += f"eval_loss={eval_loss:.4f} "
            if lr is not None:
                msg += f"lr={lr:.2e}"
            self.logger.info(msg)

    def on_evaluate(self, args, state, control, metrics=None, **kwargs) -> None:
        """On evaluate."""
        if metrics:
            self.logger.info("=== Evaluation ===")
            for k, v in metrics.items():
                if isinstance(v, (int, float)):
                    self.logger.info(f"  {k}: {v:.4f}")
                else:
                    self.logger.info(f"  {k}: {v}")
            self.tracker.log(metrics, step=state.global_step)
            self.tracker.flush()

    def on_train_end(self, args, state, control, **kwargs) -> None:
        """On train end."""
        self.tracker.save()
        self.logger.info("Training finished. Metrics saved.")

Automatically captures training/evaluation logs and stores them in a MetricsTracker.

Usage -----= tracker = MetricsTracker("metrics.json") trainer.add_callback(LoggingCallback(tracker, logger))

Ancestors

• transformers.trainer_callback.TrainerCallback

Methods

def on_evaluate(self, args, state, control, metrics=None, **kwargs) ‑> None

Expand source code

def on_evaluate(self, args, state, control, metrics=None, **kwargs) -> None:
    """On evaluate."""
    if metrics:
        self.logger.info("=== Evaluation ===")
        for k, v in metrics.items():
            if isinstance(v, (int, float)):
                self.logger.info(f"  {k}: {v:.4f}")
            else:
                self.logger.info(f"  {k}: {v}")
        self.tracker.log(metrics, step=state.global_step)
        self.tracker.flush()

On evaluate.

def on_log(self, args, state, control, logs=None, **kwargs) ‑> None

Expand source code

def on_log(self, args, state, control, logs=None, **kwargs) -> None:
    """On log."""
    if logs:
        step = state.global_step
        self.tracker.log(logs, step=step)
        # Log key metrics to the console/file logger
        loss = logs.get("loss", None)
        eval_loss = logs.get("eval_loss", None)
        lr = logs.get("learning_rate", None)
        msg = f"Step {step}: "
        if loss is not None:
            msg += f"loss={loss:.4f} "
        if eval_loss is not None:
            msg += f"eval_loss={eval_loss:.4f} "
        if lr is not None:
            msg += f"lr={lr:.2e}"
        self.logger.info(msg)

On log.

def on_train_end(self, args, state, control, **kwargs) ‑> None

Expand source code

def on_train_end(self, args, state, control, **kwargs) -> None:
    """On train end."""
    self.tracker.save()
    self.logger.info("Training finished. Metrics saved.")

On train end.

class LoggingConfig (log_level: str = 'INFO',
log_format: str = 'structured',
log_to_console: bool = True,
log_to_file: bool = True,
log_dir: str = './logs',
experiment_name: str = 'experiment',
run_id: str | None = None,
use_timestamp: bool = True,
rank: int | None = None,
rank_specific_logs: bool = True,
use_queue_handler: bool = True,
metrics_enabled: bool = True,
metrics_flush_every: int = 10,
extra_metadata: dict[str, typing.Any] = <factory>)

Expand source code

@dataclass
class LoggingConfig:
    """Complete logging & experiment tracking configuration."""

    log_level: str = "INFO"
    log_format: str = "structured"  # "structured" (JSON lines) or "human"
    log_to_console: bool = True
    log_to_file: bool = True
    log_dir: str = "./logs"
    experiment_name: str = "experiment"
    run_id: str | None = None
    use_timestamp: bool = True
    # Rank-specific and thread-safe logging
    rank: int | None = None  # FIX: Local rank for DDP isolated file logging
    rank_specific_logs: bool = True  # FIX: Enforce isolated rank_{i}.log files per process
    use_queue_handler: bool = True  # FIX: Enforce thread-safe asynchronous logging via QueueHandler
    # Metrics tracking
    metrics_enabled: bool = True
    metrics_flush_every: int = 10  # flush metrics file every N logs
    # Extra metadata added to every event
    extra_metadata: dict[str, Any] = field(default_factory=dict)

Complete logging & experiment tracking configuration.

Instance variables

var experiment_name : str

The type of the None singleton.

var extra_metadata : dict[str, typing.Any]

The type of the None singleton.

var log_dir : str

The type of the None singleton.

var log_format : str

The type of the None singleton.

var log_level : str

The type of the None singleton.

var log_to_console : bool

The type of the None singleton.

var log_to_file : bool

The type of the None singleton.

var metrics_enabled : bool

The type of the None singleton.

var metrics_flush_every : int

The type of the None singleton.

var rank : int | None

The type of the None singleton.

var rank_specific_logs : bool

The type of the None singleton.

var run_id : str | None

The type of the None singleton.

var use_queue_handler : bool

The type of the None singleton.

var use_timestamp : bool

The type of the None singleton.

class MetricsTracker (save_path: str | None = None)

Expand source code

class MetricsTracker:
    """Collects metrics during training and saves them as JSON."""

    def __init__(self, save_path: str | None = None) -> None:
        self.save_path = save_path
        self.records: list[dict[str, Any]] = []

    def log(self, metrics: dict[str, Any], step: int = None) -> None:
        """Log."""
        entry = dict(metrics)
        if step is not None:
            entry["step"] = step
        self.records.append(entry)

    def to_dataframe(self) -> Any:
        """To dataframe."""
        import pandas as pd

        return pd.DataFrame(self.records)

    def save(self, path: str | None = None) -> None:
        """Save."""
        p = path or self.save_path
        if p:
            os.makedirs(os.path.dirname(p), exist_ok=True)
            with open(p, "w") as f:
                json.dump(self.records, f, indent=2)

    def flush(self) -> None:
        """Flush."""
        if self.save_path:
            self.save(self.save_path)

Collects metrics during training and saves them as JSON.

Methods

def flush(self) ‑> None

Expand source code

def flush(self) -> None:
    """Flush."""
    if self.save_path:
        self.save(self.save_path)

Flush.

def log(self, metrics: dict[str, typing.Any], step: int = None) ‑> None

Expand source code

def log(self, metrics: dict[str, Any], step: int = None) -> None:
    """Log."""
    entry = dict(metrics)
    if step is not None:
        entry["step"] = step
    self.records.append(entry)

Log.

def save(self, path: str | None = None) ‑> None

Expand source code

def save(self, path: str | None = None) -> None:
    """Save."""
    p = path or self.save_path
    if p:
        os.makedirs(os.path.dirname(p), exist_ok=True)
        with open(p, "w") as f:
            json.dump(self.records, f, indent=2)

Save.

def to_dataframe(self) ‑> Any

Expand source code

def to_dataframe(self) -> Any:
    """To dataframe."""
    import pandas as pd

    return pd.DataFrame(self.records)

To dataframe.

class ResultsManager (experiment_name: str,
base_dir: str = 'results',
config: dict[str, typing.Any] | None = None,
log_to_file: bool = True)

Expand source code

class ResultsManager:
    """Manages per‑run result directories, logging, metrics, and trainer wrapping."""

    def __init__(
        self,
        experiment_name: str,
        base_dir: str = "results",
        config: dict[str, Any] | None = None,
        log_to_file: bool = True,
    ) -> None:
        self.experiment_name = experiment_name
        self.base_dir = Path(base_dir)
        self.config = config or {}
        self.log_to_file = log_to_file

        timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
        self.run_dir = self.base_dir / f"{experiment_name}_{timestamp}"
        self.run_dir.mkdir(parents=True, exist_ok=True)

        self.checkpoint_dir = self.run_dir / "checkpoints"
        self.log_dir = self.run_dir / "logs"
        self.final_model_dir = self.run_dir / "final_model"
        for d in [self.checkpoint_dir, self.log_dir, self.final_model_dir]:
            d.mkdir(parents=True, exist_ok=True)

        self.logger = self._setup_logger()
        self.metrics_path = self.run_dir / "metrics.jsonl"
        self.history: list[dict[str, Any]] = []

        if config:
            with open(self.run_dir / "config.yaml", "w") as f:
                yaml.dump(config, f, default_flow_style=False)

        self.logger.info(f"Results directory: {self.run_dir}")

    def _setup_logger(self) -> logging.Logger:
        """setup logger."""
        logger = logging.getLogger(f"training.{self.experiment_name}")
        logger.setLevel(logging.INFO)
        logger.handlers.clear()
        ch = logging.StreamHandler()
        ch.setFormatter(logging.Formatter("%(asctime)s - %(levelname)s - %(message)s"))
        logger.addHandler(ch)
        if self.log_to_file:
            fh = logging.FileHandler(self.log_dir / "run.log")
            fh.setFormatter(
                logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
            )
            logger.addHandler(fh)
        return logger

    def log_metrics(self, metrics: dict[str, Any], step: int | None = None) -> None:
        """Log metrics."""
        record = {"step": step} if step is not None else {}
        record.update(metrics)
        record["timestamp"] = datetime.now().isoformat()
        record = _sanitize_json_value(record)
        self.history.append(record)
        with open(self.metrics_path, "a") as f:
            f.write(json.dumps(record, allow_nan=False) + "\n")

    def save_plots(self) -> None:
        """Persist plots to disk."""
        if self.history:
            plot_path = self.run_dir / "training_history.png"
            plot_training_curves(self.history, str(plot_path))
            self.logger.info(f"Saved training plot to {plot_path}")

    def save_final_model(self, model, tokenizer=None) -> None:
        """Persist final model to disk."""
        model.save_pretrained(self.final_model_dir)
        if tokenizer:
            tokenizer.save_pretrained(self.final_model_dir)
        self.logger.info(f"Final model saved to {self.final_model_dir}")

    def get_checkpoint_dir(self) -> str:
        """Return the checkpoint dir."""
        return str(self.checkpoint_dir)

    def push_to_hf(self, repo_id: str | None = None) -> None:
        """Push full run results (models, plots, metrics) to HuggingFace Hub."""
        from neuralops.results_saver.backends import huggingface as hf_backend

        cfg = {"hf_options": {"repo_id": repo_id}} if repo_id else self.config
        hf_backend.push(self.experiment_name, self.run_dir, cfg)

    def generate_summary_report(self, final_metrics: dict[str, Any] | None = None) -> None:
        """Generate structured markdown and json summary reports for the run."""

        def _make_serializable(obj):
            if hasattr(obj, "to_dict") and callable(obj.to_dict):
                return obj.to_dict()
            if hasattr(obj, "__dict__"):
                return {
                    k: _make_serializable(v)
                    for k, v in obj.__dict__.items()
                    if not k.startswith("_")
                }
            if isinstance(obj, dict):
                return {str(k): _make_serializable(v) for k, v in obj.items()}
            if isinstance(obj, (list, tuple, set)):
                return [_make_serializable(item) for item in obj]
            if isinstance(obj, (datetime, Path)):
                return str(obj)
            if hasattr(obj, "value"):
                return obj.value
            try:
                json.dumps(obj)
                return obj
            except Exception:
                return str(obj)

        summary = {
            "experiment_name": self.experiment_name,
            "run_dir": str(self.run_dir),
            "total_steps": len(self.history),
            "config": _make_serializable(self.config),
            "final_metrics": _sanitize_json_value(
                _make_serializable(final_metrics or (self.history[-1] if self.history else {}))
            ),
            "timestamp": datetime.now().isoformat(),
        }
        summary_path = self.run_dir / "run_summary.json"
        with open(summary_path, "w") as f:
            json.dump(summary, f, indent=2, allow_nan=False)

        card_path = self.run_dir / "run_summary.md"
        with open(card_path, "w") as f:
            f.write(f"# Run Summary: {self.experiment_name}\n\n")
            f.write(f"- **Timestamp**: {summary['timestamp']}\n")
            f.write(f"- **Total Logging Steps**: {summary['total_steps']}\n")
            f.write(f"- **Run Directory**: `{self.run_dir}`\n\n")
            f.write("## Final Metrics\n```json\n")
            f.write(json.dumps(summary["final_metrics"], indent=2))
            f.write("\n```\n")
        self.logger.info(f"Generated structured run summary report at {card_path}")

    def wrap_trainer(self, trainer) -> Any:
        """Add standard callbacks and metrics logging to a Hugging Face Trainer."""
        try:
            trainer.add_callback(MemoryCallback())
            trainer.add_callback(TimeCallback())
            trainer.add_callback(GradientNormCallback(log_every_n_steps=100))
        except Exception as e:
            self.logger.warning(f"Could not add callbacks: {e}")

        original_log = trainer.log

        def custom_log(logs) -> None:
            """Custom log."""
            original_log(logs)
            self.log_metrics(logs)

        trainer.log = custom_log
        return trainer

Manages per‑run result directories, logging, metrics, and trainer wrapping.

Methods

def generate_summary_report(self, final_metrics: dict[str, typing.Any] | None = None) ‑> None

Expand source code

def generate_summary_report(self, final_metrics: dict[str, Any] | None = None) -> None:
    """Generate structured markdown and json summary reports for the run."""

    def _make_serializable(obj):
        if hasattr(obj, "to_dict") and callable(obj.to_dict):
            return obj.to_dict()
        if hasattr(obj, "__dict__"):
            return {
                k: _make_serializable(v)
                for k, v in obj.__dict__.items()
                if not k.startswith("_")
            }
        if isinstance(obj, dict):
            return {str(k): _make_serializable(v) for k, v in obj.items()}
        if isinstance(obj, (list, tuple, set)):
            return [_make_serializable(item) for item in obj]
        if isinstance(obj, (datetime, Path)):
            return str(obj)
        if hasattr(obj, "value"):
            return obj.value
        try:
            json.dumps(obj)
            return obj
        except Exception:
            return str(obj)

    summary = {
        "experiment_name": self.experiment_name,
        "run_dir": str(self.run_dir),
        "total_steps": len(self.history),
        "config": _make_serializable(self.config),
        "final_metrics": _sanitize_json_value(
            _make_serializable(final_metrics or (self.history[-1] if self.history else {}))
        ),
        "timestamp": datetime.now().isoformat(),
    }
    summary_path = self.run_dir / "run_summary.json"
    with open(summary_path, "w") as f:
        json.dump(summary, f, indent=2, allow_nan=False)

    card_path = self.run_dir / "run_summary.md"
    with open(card_path, "w") as f:
        f.write(f"# Run Summary: {self.experiment_name}\n\n")
        f.write(f"- **Timestamp**: {summary['timestamp']}\n")
        f.write(f"- **Total Logging Steps**: {summary['total_steps']}\n")
        f.write(f"- **Run Directory**: `{self.run_dir}`\n\n")
        f.write("## Final Metrics\n```json\n")
        f.write(json.dumps(summary["final_metrics"], indent=2))
        f.write("\n```\n")
    self.logger.info(f"Generated structured run summary report at {card_path}")

Generate structured markdown and json summary reports for the run.

def get_checkpoint_dir(self) ‑> str

Expand source code

def get_checkpoint_dir(self) -> str:
    """Return the checkpoint dir."""
    return str(self.checkpoint_dir)

Return the checkpoint dir.

def log_metrics(self, metrics: dict[str, typing.Any], step: int | None = None) ‑> None

Expand source code

def log_metrics(self, metrics: dict[str, Any], step: int | None = None) -> None:
    """Log metrics."""
    record = {"step": step} if step is not None else {}
    record.update(metrics)
    record["timestamp"] = datetime.now().isoformat()
    record = _sanitize_json_value(record)
    self.history.append(record)
    with open(self.metrics_path, "a") as f:
        f.write(json.dumps(record, allow_nan=False) + "\n")

Log metrics.

def push_to_hf(self, repo_id: str | None = None) ‑> None

Expand source code

def push_to_hf(self, repo_id: str | None = None) -> None:
    """Push full run results (models, plots, metrics) to HuggingFace Hub."""
    from neuralops.results_saver.backends import huggingface as hf_backend

    cfg = {"hf_options": {"repo_id": repo_id}} if repo_id else self.config
    hf_backend.push(self.experiment_name, self.run_dir, cfg)

Push full run results (models, plots, metrics) to HuggingFace Hub.

def save_final_model(self, model, tokenizer=None) ‑> None

Expand source code

def save_final_model(self, model, tokenizer=None) -> None:
    """Persist final model to disk."""
    model.save_pretrained(self.final_model_dir)
    if tokenizer:
        tokenizer.save_pretrained(self.final_model_dir)
    self.logger.info(f"Final model saved to {self.final_model_dir}")

Persist final model to disk.

def save_plots(self) ‑> None

Expand source code

def save_plots(self) -> None:
    """Persist plots to disk."""
    if self.history:
        plot_path = self.run_dir / "training_history.png"
        plot_training_curves(self.history, str(plot_path))
        self.logger.info(f"Saved training plot to {plot_path}")

Persist plots to disk.

def wrap_trainer(self, trainer) ‑> Any

Expand source code

def wrap_trainer(self, trainer) -> Any:
    """Add standard callbacks and metrics logging to a Hugging Face Trainer."""
    try:
        trainer.add_callback(MemoryCallback())
        trainer.add_callback(TimeCallback())
        trainer.add_callback(GradientNormCallback(log_every_n_steps=100))
    except Exception as e:
        self.logger.warning(f"Could not add callbacks: {e}")

    original_log = trainer.log

    def custom_log(logs) -> None:
        """Custom log."""
        original_log(logs)
        self.log_metrics(logs)

    trainer.log = custom_log
    return trainer

Add standard callbacks and metrics logging to a Hugging Face Trainer.