Raise a cumstom exception in on_fit_end

[nzw0301]

Hi

I would like to raise an exception in the on_fit_end of a callback with DDP training. For example, the following code calls optuna.exceptions.TrialPruned() in the method.

from typing import Dict
from typing import List

import optuna
from pytorch_lightning.callbacks import Callback
import pytorch_lightning as pl
import torch
from torch import nn
import torch.nn.functional as F


class Model(pl.LightningModule):
    def __init__(self) -> None:

        super().__init__()
        self._model = nn.Sequential(nn.Linear(4, 8))

    def forward(self, data: torch.Tensor) -> torch.Tensor:  # type: ignore

        return self._model(data)

    def training_step(  # type: ignore
        self, batch: List[torch.Tensor], batch_nb: int
    ) -> Dict[str, torch.Tensor]:

        data, target = batch
        output = self.forward(data)
        loss = F.nll_loss(output, target)
        return {"loss": loss}

    def validation_step(  # type: ignore
        self, batch: List[torch.Tensor], batch_nb: int
    ) -> Dict[str, torch.Tensor]:

        data, target = batch
        output = self.forward(data)
        pred = output.argmax(dim=1, keepdim=True)
        accuracy = pred.eq(target.view_as(pred)).double().mean()
        return {"validation_accuracy": accuracy}

    def validation_epoch_end(self, outputs: List[Dict[str, torch.Tensor]]) -> None:
        if not len(outputs):
            return

        accuracy = sum(x["validation_accuracy"] for x in outputs) / len(outputs)
        self.log("accuracy", accuracy)

    def configure_optimizers(self) -> torch.optim.Optimizer:

        return torch.optim.SGD(self._model.parameters(), lr=1e-2)

    def train_dataloader(self) -> torch.utils.data.DataLoader:

        return self._generate_dummy_dataset()

    def val_dataloader(self) -> torch.utils.data.DataLoader:

        return self._generate_dummy_dataset()

    def _generate_dummy_dataset(self) -> torch.utils.data.DataLoader:

        data = torch.zeros(3, 4, dtype=torch.float32)
        target = torch.zeros(3, dtype=torch.int64)
        dataset = torch.utils.data.TensorDataset(data, target)
        return torch.utils.data.DataLoader(dataset, batch_size=1)


class ModelDDP(Model):
    def __init__(self) -> None:

        super().__init__()

    def validation_step(  # type: ignore
        self, batch: List[torch.Tensor], batch_nb: int
    ) -> Dict[str, torch.Tensor]:

        data, target = batch
        output = self.forward(data)
        pred = output.argmax(dim=1, keepdim=True)
        accuracy = pred.eq(target.view_as(pred)).double().mean()

        if self.local_rank == 0:
            accuracy = torch.tensor(0.3)
        elif self.local_rank == 1:
            accuracy = torch.tensor(0.6)

        self.log("accuracy", accuracy, sync_dist=True)


class CustomCallback(Callback):
    def __init__(self) -> None:
        super().__init__()

    def on_fit_end(self, trainer, pl_module) -> None:
        raise optuna.exceptions.TrialPruned()

if __name__ == '__main__':


    trainer = pl.Trainer(
        max_epochs=2,
        accelerator="ddp_cpu",
        num_processes=2,
        enable_checkpointing=False,
        callbacks=[CustomCallback()],
    )

    model = ModelDDP()
    trainer.fit(model)

If I use pl 1.5.*, on_fit_end is called once. So I could handle the exception.

However, if I use pl 1.6.0, on_fit_end is called by every child process. As a result, we got torch.multiprocessing.spawn.ProcessRaisedException due to optuna.exceptions.TrialPruned().

How should I raise a custom exception as with pl 1.5.*when we use pl 1.6.0 without torch.multiprocessing.spawn.ProcessRaisedException?

Note that my env is as follows:

• OS: the latest M1 mac
• PyTorch: 1.11.0

Best regards,

[nzw0301]

More specifically, it is used in a callback developed by Optuna for PyTorch-Lightning.

CC: @Borda @akihironitta because I saw your reactions to an issue.

[akihironitta]

Hi @nzw0301! Sorry for replying very late.

I'm sure the change in the behaviour was introduced in #10896.

@awaelchli Could I ask for your help here? Is it possible for a user to catch an exception of their interest instead of torch.multiprocessing.spawn.ProcessRaisedException raised from the hook when using spawn-based strategies?

[awaelchli]

Yes, it is not possible to obtain the original exception class, a limitation of the spawn method (ddp_cpu, ddp_spawn, and tpu). We decided to run on_fit_end in the worker process for consistency with on_fit_start who also runs in the worker process.
We currently don't have any hook that runs at the end of fit() in the main process only.

What you could do in your callback is implement the on_exception() hook. There, you get the exception as input. If it is a ProcessRaisedException and the message matches the one from optuna.exceptions.TrialPruned(), then re-create that exception with the proper class and re-raise it. Let me know if that helps.

Unfortunately, I don't fully understand why raising this special exception is needed in the first place. It must have something to to with a special control flow around the fit call. Do all processes raise that, or only some of them?

[nzw0301]

@awaelchli Thank you for your informative response! I'll check on_exception!

Do all processes raise that, or only some of them?

In our case, one process is enough, so we originally raised optuna.exceptions.TrialPruned() on only the main process. As you suggested, we could raise it in the fit, but ideally, we wanted to raise it via callback for consistency with our API design.