# Callbacks

All built-in callbacks and the `FabricCallbackProtocol` hook interface.

Source: `common/callbacks.py`

## FabricCallbackProtocol

The protocol that all callbacks implement. Each method is optional — implement only the hooks you need.

### Hook signatures

```python
class FabricCallbackProtocol(Protocol):
    # Training lifecycle
    def on_train_start(self, *, fabric, model, optimizer, lr_scheduler, dataloader, hparams, **kwargs): ...
    def on_train_end(self, **kwargs): ...
    def on_test_start(self, *, dataloader, **kwargs): ...
    def on_test_end(self, *, metrics, **kwargs): ...

    # Epoch level
    def on_epoch_start(self, epoch: int, max_epochs: int, **kwargs): ...
    def on_epoch_end(self, epoch: int, *, train_metrics, val_metrics, **kwargs): ...
    def on_train_epoch_start(self, epoch: int, max_epochs: int, **kwargs): ...
    def on_train_epoch_end(self, epoch: int, *, metrics, **kwargs): ...
    def on_validation_epoch_start(self, epoch: int, *, epoch_interval, **kwargs): ...
    def on_validation_epoch_end(self, epoch: int, *, metrics, **kwargs): ...

    # Batch level
    def on_train_batch_start(self, batch, batch_idx: int, **kwargs): ...
    def on_train_batch_end(self, batch, batch_idx: int, *, metrics, **kwargs): ...
    def on_validation_batch_start(self, batch, batch_idx: int, **kwargs): ...
    def on_validation_batch_end(self, batch, batch_idx: int, *, metrics, **kwargs): ...
    def on_test_batch_start(self, batch, batch_idx: int, **kwargs): ...
    def on_test_batch_end(self, batch, batch_idx: int, *, metrics, **kwargs): ...
```

### Hook execution order per epoch

```
on_epoch_start(epoch, max_epochs)
  on_train_epoch_start(epoch, max_epochs)
    [for each batch]:
      on_train_batch_start(batch, batch_idx)
      ... training_step ...
      on_train_batch_end(batch, batch_idx, metrics=step_metrics)
  on_train_epoch_end(epoch, metrics=train_epoch_metrics)

  on_validation_epoch_start(epoch, epoch_interval)
    [for each batch]:
      on_validation_batch_start(batch, batch_idx)
      ... validation_step ...
      on_validation_batch_end(batch, batch_idx, metrics=step_metrics)
  on_validation_epoch_end(epoch, metrics=val_epoch_metrics)
on_epoch_end(epoch, train_metrics=..., val_metrics=...)
```

### Registering callbacks

Via `default_fabric`:
```python
fabric = default_fabric(fn_custom_callback=lambda v: [MyCallback()])
```

Direct append (less common):
```python
fabric._callbacks.append(MyCallback())
```

### Finding callbacks

```python
from zyplib.train.common.callbacks import find_callback, SaveCheckpoints
ckpt_cb = find_callback(fabric, SaveCheckpoints)  # returns instance or None
```

## SaveCheckpoints

Automatic model checkpoint saving with Top-K best model strategy.

### Constructor

```python
SaveCheckpoints(
    dir: str,                                          # checkpoint directory
    save_every_k: int | None = None,                   # periodic save interval
    save_best: bool = True,                            # save best models
    keep_best_k: int | None = 1,                       # keep top-K best
    monitor: str = 'loss',                             # metric to monitor
    monitor_on: Literal['train', 'val'] = 'val',       # which phase
    less_is_better: bool = True,                       # lower is better?
    ckpt_ext: str = '.ckpt',                           # file extension
    if_not_empty: Literal['ask', 'clean', 'error', 'ignore'] = 'ask',
)
```

### Behavior

- Listens to `on_train_epoch_end` and `on_validation_epoch_end` to collect metrics
- On `on_epoch_end`, evaluates whether to save based on Top-K strategy
- `best_checkpoint()` returns the path to the best checkpoint file
- `best_score` property returns the best monitored metric value

### Monitor metric resolution

Looks up metric in this order:
1. `metrics[monitor]` (exact match)
2. `metrics[f'val_{monitor}']` (with val prefix)
3. `metrics[f'train_{monitor}']` (with train prefix)

### Note on if_not_empty

- `'ask'`: interactive prompt — **blocks in non-interactive environments**
- `'clean'`: silently removes existing checkpoints
- `'error'`: raises exception
- `'ignore'`: continues using existing directory

Usually configured via `default_fabric` parameters (`save_best`, `keep_best_k`, `monitor`, etc.) rather than directly.

## PrintTrainInfo

Prints formatted training progress to console and optionally to file.

### Constructor

```python
PrintTrainInfo(out_dir: str | None = None)
```

When `out_dir` is provided, also writes to `{out_dir}/train-info.txt`.

### Hooks used

- `on_train_start`: prints start time
- `on_train_end`: prints end time and elapsed duration
- `on_train_epoch_start`: prints epoch separator
- `on_train_epoch_end`: prints train metrics
- `on_validation_epoch_end`: prints val metrics
- `on_test_end`: prints test metrics

## SaveHparams

Saves hyperparameters to `hparams.yaml` and TensorBoard.

### Constructor

```python
SaveHparams(dir: str, inject_namespace: str = 'default')
```

### Hooks used

- `on_train_start`: saves hparams from kwargs or `simple_inject.state(namespace)`

### Hparams sources (priority order)

1. `hparams` kwarg passed to `train_loop` / `trainer.fit`
2. `simple_inject.state(inject_namespace)` dict
3. If neither available, prints warning

## EarlyStopping

Stops training when monitored metric stops improving.

### Constructor

```python
EarlyStopping(
    monitor: str = 'loss',
    monitor_on: Literal['train', 'val'] = 'val',
    min_delta: float = 0.0,
    patience: int = 10,
    mode: Literal['min', 'max'] = 'min',
    verbose: bool = True,
    fn_update_callback: ((counter, patience) -> None) | None = None,
    fn_stop_callback: ((epoch) -> None) | None = None,
)
```

### Behavior

- Monitors metric on `on_validation_epoch_end` (if `monitor_on='val'`) or `on_train_epoch_end` (if `monitor_on='train'`)
- Sets `stop_training = True` after `patience` epochs without improvement
- `train_loop` checks `stop_training` flag at the start of each epoch

### Properties

- `best_score`: best observed metric value
- `best_epoch`: epoch of best score
- `counter`: current patience counter
- `stop_training`: whether training should stop

### Registration

Must be added via `fn_custom_callback` — not included in `default_fabric` by default:
```python
fabric = default_fabric(
    fn_custom_callback=lambda v: [EarlyStopping(monitor='loss', patience=20)]
)
```

## LRSchedulerCallback

Flexible learning rate scheduling with epoch and step modes.

### Constructor

```python
LRSchedulerCallback(
    lr_scheduler: _LRScheduler,
    interval: Literal['epoch', 'step', 'auto'] = 'auto',
    frequency: int = 1,
    monitor_metric: str = 'loss',
    use_val_metric: bool = True,
)
```

### Auto interval detection

When `interval='auto'`:
- `OneCycleLR`, `CyclicLR`, `CosineAnnealingWarmRestarts` -> `'step'`
- All others -> `'epoch'`

### ReduceLROnPlateau handling

Automatically detected. When the scheduler is `ReduceLROnPlateau`:
- Passes the monitored metric value to `scheduler.step(metric)`
- Uses val metrics if available and `use_val_metric=True`, otherwise falls back to train metrics

### Registration

```python
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)
fabric = default_fabric(
    fn_custom_callback=lambda v: [LRSchedulerCallback(scheduler)]
)
```

### Backward compatibility

If `lr_scheduler` is passed directly to `procedure.train_loop(...)`, a `LRSchedulerCallback` is auto-injected. This is deprecated — prefer explicit registration.

## MonitorLearningRate

Logs current learning rate to fabric loggers.

### Constructor

```python
MonitorLearningRate(fabric: Fabric | None = None)
```

### Hooks used

- `on_train_start`: captures optimizer reference
- `on_train_epoch_start`: logs LR for each param group

Logs `lr_group_0`, `lr_group_1`, etc. If single param group, also logs `learning_rate`.

## FabricLogMetrics

Automatically logs all metrics to all fabric loggers. Auto-injected when any logger is enabled.

### Hooks used

- `on_train_batch_end`: logs with prefix `train/batch/`
- `on_validation_batch_end`: logs with prefix `val/batch/`
- `on_train_epoch_end`: logs with prefix `train/epoch/` and step=epoch
- `on_validation_epoch_end`: logs with prefix `val/epoch/` and step=epoch
- `on_test_end`: logs with prefix `test/`

### Note

You rarely need to interact with this directly — it's auto-injected by `default_fabric` when loggers are present.

## SaveLoggersFinally

Flushes all loggers on training/test end. Auto-injected when any logger is enabled.

- Calls `logger.save()` for all loggers
- For TensorBoard: converts event files to JSON for easy programmatic access
