ValueError: 35 is not a valid Signals when running Scalene memory profiling on PyTorch Lightning AI

[belfaunder]

I would like to use Scalene for memory profiling in a PyTorch Lightning AI project, but I am running into a signal-handling issue during teardown.

Behavior:

When I run: scalene run --memory main.py, the program fails with ValueError: 35 is not a valid Signals during teardown phase of Lightning AI.

Full traceback:

Error in program being profiled:
 35 is not a valid Signals
Traceback (most recent call last):
  File "/home/belfaunder/test_scalene/.venv/lib/python3.12/site-packages/pytorch_lightning/trainer/call.py", line 49, in _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/belfaunder/test_scalene/.venv/lib/python3.12/site-packages/pytorch_lightning/trainer/trainer.py", line 630, in _fit_impl
    self._run(model, ckpt_path=ckpt_path, weights_only=weights_only)
  File "/home/belfaunder/test_scalene/.venv/lib/python3.12/site-packages/pytorch_lightning/trainer/trainer.py", line 1085, in _run
    self._teardown()
  File "/home/belfaunder/test_scalene/.venv/lib/python3.12/site-packages/pytorch_lightning/trainer/trainer.py", line 1108, in _teardown
    self._signal_connector.teardown()
  File "/home/belfaunder/test_scalene/.venv/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/signal_connector.py", line 129, in teardown
    self._register_signal(signum, handler)
  File "/home/belfaunder/test_scalene/.venv/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/signal_connector.py", line 153, in _register_signal
    signal.signal(signum, handlers)  # type: ignore[arg-type]
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/belfaunder/test_scalene/.venv/lib/python3.12/site-packages/scalene/replacement_signal_fns.py", line 113, in replacement_signal
    signal.Signals(signum) in expected_handlers_map
    ^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3.12/enum.py", line 757, in __call__
    return cls.__new__(cls, value)
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3.12/enum.py", line 1171, in __new__
    raise ve_exc
ValueError: 35 is not a valid Signals

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/belfaunder/test_scalene/.venv/lib/python3.12/site-packages/scalene/scalene_profiler.py", line 1269, in profile_code
    exec(code, the_globals, the_locals)
  File "/home/belfaunder/test_scalene/.venv/lib/python3.12/site-packages/scalene/replacement_exec.py", line 95, in exec_replacement
    orig_exec(__source, globals, locals, **kwargs)
  File "/home/belfaunder/test_scalene/main.py", line 49, in <module>
    trainer.fit(mnist_model, train_loader)
  File "/home/belfaunder/test_scalene/.venv/lib/python3.12/site-packages/pytorch_lightning/trainer/trainer.py", line 584, in fit
    call._call_and_handle_interrupt(
  File "/home/belfaunder/test_scalene/.venv/lib/python3.12/site-packages/pytorch_lightning/trainer/call.py", line 70, in _call_and_handle_interrupt
    trainer._teardown()
  File "/home/belfaunder/test_scalene/.venv/lib/python3.12/site-packages/pytorch_lightning/trainer/trainer.py", line 1108, in _teardown
    self._signal_connector.teardown()
  File "/home/belfaunder/test_scalene/.venv/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/signal_connector.py", line 129, in teardown
    self._register_signal(signum, handler)
  File "/home/belfaunder/test_scalene/.venv/lib/python3.12/site-packages/pytorch_lightning/trainer/connectors/signal_connector.py", line 153, in _register_signal
    signal.signal(signum, handlers)  # type: ignore[arg-type]
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/belfaunder/test_scalene/.venv/lib/python3.12/site-packages/scalene/replacement_signal_fns.py", line 113, in replacement_signal
    signal.Signals(signum) in expected_handlers_map
    ^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3.12/enum.py", line 757, in __call__
    return cls.__new__(cls, value)
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3.12/enum.py", line 1171, in __new__
    raise ve_exc
ValueError: 35 is not a valid Signals

Environment

Python: 3.12.3

lightning 2.6.1
scalene: 2.1.4
torch: 2.10.0
torchvision: 0.25.0

OS: Ubuntu 24.04 running under WSL2

Minimal reproduction example

The full code of main.py module for reproducing is based on example from Lightning AI (https://lightning.ai/docs/pytorch/stable/notebooks/lightning_examples/mnist-hello-world.html):

import os
import pytorch_lightning as pl
import torch
from torch.nn import functional as F
from torch.utils.data import DataLoader, Subset
from torchvision import transforms
from torchvision.datasets import MNIST

class MNISTModel(pl.LightningModule):
    """A PyTorch Lightning module for classifying images in the MNIST dataset."""
    def __init__(self):
        """Initializes a new instance of the MNISTModel class."""
        super().__init__()
        self.l1 = torch.nn.Linear(28 * 28, 10)
        self.leaked_losses = []

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        flattened = x.view(x.size(0), -1)
        hidden = self.l1(flattened)
        activated = torch.relu(hidden)
        return activated

    def training_step(self, batch, batch_idx: int):
        x, y = batch
        loss = F.cross_entropy(self(x), y)
        big = loss.repeat(1,1000)
        self.leaked_losses.append(big)

        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=0.02)
    
# Init the model
mnist_model = MNISTModel()

# Init DataLoader from MNIST Dataset
train_ds = MNIST(os.environ.get("PATH_DATASETS", "."), train=True, download=True, transform=transforms.ToTensor())
train_ds_small = Subset(train_ds, range(1000))

# Create a dataloader
train_loader = DataLoader(train_ds_small, batch_size=3)

# Initialize a trainer
trainer = pl.Trainer(
    accelerator="cpu",
    #strategy="ddp",
    devices=1,
    max_epochs=3,
)

# Train the model 
trainer.fit(mnist_model, train_loader)