Rich progress bar error when resume training

[YAndrewL]

Bug description

I encountered an error when resuming training with the rich progress bar. An assertion error says progress bar is None

What version are you seeing the problem on?

v2.5

Reproduced in studio

No response

How to reproduce the bug

callbacks = [
    checkpoint_callback,
    RichProgressBar(),
    lr_monitor,
    time_limit_callback,
]



if checkpoint_path:
    trainer.fit(model, data_module, ckpt_path=checkpoint_path)
else:
    trainer.fit(model, data_module)

Error messages and logs

It's an assertion error that is defined as assert progress_bar_id is not None .

Environment

PyTorch Lightning Version (e.g., 2.5.0):
PyTorch Version (e.g., 2.6.0):
Python version (e.g., 3.12.9):
OS (e.g., Linux):

More info

No response

[leoyala]

I am having the same error in pytorch-lightning v2.5.2

[bhimrazy]

Hi @YAndrewL and @leoyala — thanks for raising the issue!

I tried to reproduce it on my end but couldn’t manage to. Would you mind sharing a minimal, complete reproducible example, if possible? That would be super helpful.

[YAndrewL]

Hi @bhimrazy, thanks for your reply, here is a minimal example:

import argparse

import pytorch_lightning as pl
import torch
import torch.nn as nn
from pytorch_lightning.callbacks import ModelCheckpoint, RichProgressBar
from torch.utils.data import DataLoader, TensorDataset

# Define callbacks
checkpoint_callback = ModelCheckpoint(
    monitor="val_loss",
    dirpath="checkpoints",
    filename="model-{epoch:02d}-{val_loss:.2f}",
    save_top_k=3,
    mode="min",
)


callbacks = [
    checkpoint_callback,
    RichProgressBar(),
]


class SimpleModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 1)
        self.loss_fn = nn.MSELoss()

    def forward(self, x):
        return self.linear(x)

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = self.loss_fn(y_hat, y)
        self.log("train_loss", loss)
        return loss

    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = self.loss_fn(y_hat, y)
        self.log("val_loss", loss)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.01)


def create_random_data(num_samples=1000):
    x = torch.randn(num_samples, 10)
    y = torch.randn(num_samples, 1)
    return x, y


if __name__ == "__main__":
    # Add argument parser
    parser = argparse.ArgumentParser(description="Train a simple PyTorch Lightning model")
    parser.add_argument("--checkpoint", type=str, default=None, help="Path to checkpoint file to resume training from")

    args = parser.parse_args()

    x_train, y_train = create_random_data(1000)
    x_val, y_val = create_random_data(200)

    train_dataset = TensorDataset(x_train, y_train)
    val_dataset = TensorDataset(x_val, y_val)

    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=16)
    val_loader = DataLoader(val_dataset, batch_size=32, num_workers=16)

    model = SimpleModel()

    trainer = pl.Trainer(
        max_epochs=5,
        enable_progress_bar=True,
        log_every_n_steps=1,
        val_check_interval=0.5,
        callbacks=callbacks,
        devices=[0],
    )

    if args.checkpoint:
        trainer.fit(model, train_loader, val_loader, ckpt_path=args.checkpoint)
    else:
        trainer.fit(model, train_loader, val_loader)

and the error is shown in resume training:

lib/python3.12/site-packages/pytorch_lightning/callbacks/progress/rich_progress.py", line 452, in _update
    assert progress_bar_id is not None
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError

[bhimrazy]

Thank you, @YAndrewL. I was able to reproduce the issue when resuming from the checkpoint.

Logs

⚡ ~ python main.py --checkpoint checkpoints/model-epoch=03-val_loss=1.05.ckpt
/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/utils/data/dataloader.py:626: UserWarning: This DataLoader will create 8 worker processes in total. Our suggested max number of worker in current system is 4, which is smaller than what this DataLoader is going to create. Please be aware that excessive worker creation might get DataLoader running slow or even freeze, lower the worker number to avoid potential slowness/freeze if necessary.
  warnings.warn(
GPU available: False, used: False
TPU available: False, using: 0 TPU cores
HPU available: False, using: 0 HPUs
/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/pytorch_lightning/callbacks/model_checkpoint.py:658: Checkpoint directory /teamspace/studios/this_studio/checkpoints exists and is not empty.
Restoring states from the checkpoint path at checkpoints/model-epoch=03-val_loss=1.05.ckpt
┏━━━┳━━━━━━━━━┳━━━━━━━━━┳━━━━━━━━┳━━━━━━━┓
┃   ┃ Name    ┃ Type    ┃ Params ┃ Mode  ┃
┡━━━╇━━━━━━━━━╇━━━━━━━━━╇━━━━━━━━╇━━━━━━━┩
│ 0 │ linear  │ Linear  │     11 │ train │
│ 1 │ loss_fn │ MSELoss │      0 │ train │
└───┴─────────┴─────────┴────────┴───────┘
Trainable params: 11                                                                                                                           
Non-trainable params: 0                                                                                                                        
Total params: 11                                                                                                                               
Total estimated model params size (MB): 0                                                                                                      
Modules in train mode: 2                                                                                                                       
Modules in eval mode: 0                                                                                                                        
Restored all states from the checkpoint at checkpoints/model-epoch=03-val_loss=1.05.ckpt
/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/utils/data/dataloader.py:626: UserWarning: This DataLoader will create 8 worker processes in total. Our suggested max number of worker in current system is 4, which is smaller than what this DataLoader is going to create. Please be aware that excessive worker creation might get DataLoader running slow or even freeze, lower the worker number to avoid potential slowness/freeze if necessary.
  warnings.warn(
/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/pytorch_lightning/loops/training_epoch_loop.py:223: You're resuming from a 
checkpoint that ended before the epoch ended and your dataloader is not resumable. This can cause unreliable results if further training is 
done. Consider using an end-of-epoch checkpoint or make your dataloader resumable by implementing the `state_dict` / `load_state_dict` 
interface.

Traceback (most recent call last):
  File "/teamspace/studios/this_studio/main.py", line 92, in <module>
    trainer.fit(model, train_loader, val_loader, ckpt_path=args.checkpoint)
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 561, in fit
    call._call_and_handle_interrupt(
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 48, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 599, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 1012, in _run
    results = self._run_stage()
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 1056, in _run_stage
    self.fit_loop.run()
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/pytorch_lightning/loops/fit_loop.py", line 216, in run
    self.advance()
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/pytorch_lightning/loops/fit_loop.py", line 455, in advance
    self.epoch_loop.run(self._data_fetcher)
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/pytorch_lightning/loops/training_epoch_loop.py", line 152, in run
    self.advance(data_fetcher)
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/pytorch_lightning/loops/training_epoch_loop.py", line 363, in advance
    call._call_callback_hooks(trainer, "on_train_batch_end", batch_output, batch, batch_idx)
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 227, in _call_callback_hooks
    fn(trainer, trainer.lightning_module, *args, **kwargs)
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/pytorch_lightning/callbacks/progress/rich_progress.py", line 534, in on_train_batch_end
    self._update(self.train_progress_bar_id, batch_idx + 1)
  File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/pytorch_lightning/callbacks/progress/rich_progress.py", line 452, in _update
    assert progress_bar_id is not None
AssertionError
⚡ ~ 

[YAndrewL]

Yes, and it happens no matter multi or single GPU

[destroy314]

It seems that the issue is caused by RichProgressBar.train_progress_bar_id not being properly initialized. Skipping this update at site-packages/lightning/pytorch/callbacks/progress/rich_progress.py:450 appears to resolve the problem, though it's unclear whether this workaround has any side effects.

def _update(self, progress_bar_id: Optional["TaskID"], current: int, visible: bool = True) -> None:
    if self.progress is not None and self.is_enabled:
+       if progress_bar_id is None: return
        assert progress_bar_id is not None
        total = self.progress.tasks[progress_bar_id].total
        ...

[deependujha]

hey @YAndrewL @leoyala @destroy314 @bhimrazy

I've pushed a fix for the issue. Would appreciate it if you could take a look when you get a chance.

Thanks!
Deependu