Comet logger does not support "after_save_checkpoint" method?

[hieubnt235]

Description & Motivation

Hi, I wonder why CometLogger does not support "after_save_checkpoint" method? It's just not in your plan for now, or there something with Comet that make this method can not be supported like neptune or wandb one?

Thank you.

Pitch

No response

Alternatives

No response

Additional context

No response

cc @lantiga @Borda

[SkafteNicki]

Hi @hieubnt235, thanks for raising this issue. CometLogger does support after_save_checkpoint, but it does not contain any logic for CometLogger compared to the other loggers you mention which is the reason why it is not mention in the docs. But every hook of the base Logger class is supported by all loggers.

Here is a small example to showcase the hook being called:

import torch
import torch.nn as nn
from lightning.pytorch import LightningModule, Trainer
from lightning.pytorch.loggers import CometLogger
from lightning.pytorch.callbacks import ModelCheckpoint


class CustomCometLogger(CometLogger):
    """CometLogger subclass that proves after_save_checkpoint works"""
    
    def after_save_checkpoint(self, checkpoint_callback):
        print("🎉 after_save_checkpoint method called on CometLogger!")
        print(f"   Checkpoint callback: {type(checkpoint_callback).__name__}")
        if hasattr(checkpoint_callback, 'last_model_path'):
            print(f"   Last model path: {checkpoint_callback.last_model_path}")
        # Call parent implementation
        return super().after_save_checkpoint(checkpoint_callback)


class SimpleModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = nn.Linear(10, 1)
        
    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.layer(x)
        loss = nn.functional.mse_loss(y_hat, y)
        self.log('train_loss', loss)
        return loss
    
    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.02)


def create_dummy_data():
    """Create simple dummy data for training"""
    data = []
    for _ in range(40):  # 10 batches of size 4
        x = torch.randn(10)
        y = torch.randn(1)
        data.append((x, y))
    
    dataset = torch.utils.data.TensorDataset(
        torch.stack([d[0] for d in data]),
        torch.stack([d[1] for d in data])
    )
    return torch.utils.data.DataLoader(dataset, batch_size=4)


if __name__ == "__main__":
    # Create our custom logger (uses real CometML)
    logger = CustomCometLogger(
        project="pytorch-lightning-demo",
        experiment_name="after_save_checkpoint_demo"
    )
    
    # Create model and data
    model = SimpleModel()
    dataloader = create_dummy_data()
    
    # Create checkpoint callback that will trigger after_save_checkpoint
    checkpoint_callback = ModelCheckpoint(
        filename='demo-{epoch:02d}-{train_loss:.2f}',
        save_top_k=1,
        monitor='train_loss',
        mode='min',
        save_last=True,
        every_n_epochs=1
    )
    
    # Create trainer
    trainer = Trainer(
        logger=logger,
        callbacks=[checkpoint_callback],
        max_epochs=2,
        limit_train_batches=10,
        enable_progress_bar=False,
        enable_model_summary=False,
        log_every_n_steps=5
    )
    
    # Train the model - this will trigger after_save_checkpoint
    trainer.fit(model, dataloader)

Closing issue.

[hieubnt235]

Hi @hieubnt235, thanks for raising this issue. CometLogger does support after_save_checkpoint, but it does not contain any logic for CometLogger compared to the other loggers you mention which is the reason why it is not mention in the docs. But every hook of the base Logger class is supported by all loggers.

Here is a small example to showcase the hook being called:

import torch
import torch.nn as nn
from lightning.pytorch import LightningModule, Trainer
from lightning.pytorch.loggers import CometLogger
from lightning.pytorch.callbacks import ModelCheckpoint


class CustomCometLogger(CometLogger):
    """CometLogger subclass that proves after_save_checkpoint works"""
    
    def after_save_checkpoint(self, checkpoint_callback):
        print("🎉 after_save_checkpoint method called on CometLogger!")
        print(f"   Checkpoint callback: {type(checkpoint_callback).__name__}")
        if hasattr(checkpoint_callback, 'last_model_path'):
            print(f"   Last model path: {checkpoint_callback.last_model_path}")
        # Call parent implementation
        return super().after_save_checkpoint(checkpoint_callback)


class SimpleModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = nn.Linear(10, 1)
        
    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.layer(x)
        loss = nn.functional.mse_loss(y_hat, y)
        self.log('train_loss', loss)
        return loss
    
    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.02)


def create_dummy_data():
    """Create simple dummy data for training"""
    data = []
    for _ in range(40):  # 10 batches of size 4
        x = torch.randn(10)
        y = torch.randn(1)
        data.append((x, y))
    
    dataset = torch.utils.data.TensorDataset(
        torch.stack([d[0] for d in data]),
        torch.stack([d[1] for d in data])
    )
    return torch.utils.data.DataLoader(dataset, batch_size=4)


if __name__ == "__main__":
    # Create our custom logger (uses real CometML)
    logger = CustomCometLogger(
        project="pytorch-lightning-demo",
        experiment_name="after_save_checkpoint_demo"
    )
    
    # Create model and data
    model = SimpleModel()
    dataloader = create_dummy_data()
    
    # Create checkpoint callback that will trigger after_save_checkpoint
    checkpoint_callback = ModelCheckpoint(
        filename='demo-{epoch:02d}-{train_loss:.2f}',
        save_top_k=1,
        monitor='train_loss',
        mode='min',
        save_last=True,
        every_n_epochs=1
    )
    
    # Create trainer
    trainer = Trainer(
        logger=logger,
        callbacks=[checkpoint_callback],
        max_epochs=2,
        limit_train_batches=10,
        enable_progress_bar=False,
        enable_model_summary=False,
        log_every_n_steps=5
    )
    
    # Train the model - this will trigger after_save_checkpoint
    trainer.fit(model, dataloader)

Closing issue.

Thanks for your response, but I know that it has that method, the actual question is why this method is kept blank and not implemented like the others in the case of Comet. Actually i can custom it, but i just can't ensure this is not implemented because out of plan or it's intentionally kept blank because of some conflict.

[SkafteNicki]

It is intentionally left blank because there was no need to have any logic in that method for comet logger. For the other loggers the method has an use. If you have a specific use case where for comet where including logic in this method would be the way forward, then we can reopen the issue.