DDP spawn no longer works in Jupyter environment

[Borda]

🐛 Bug

Running multi-GPU training does not update training progress bar, used in Grid Jupyterlab

Please reproduce using the BoringModel

https://github.com/Borda/kaggle_plant-pathology/blob/main/notebooks/Plant-Pathology-with-Lightning.ipynb

To Reproduce

just change

    gpus=2
    accelerator="ddp_spawn"
    precision=16,
    accumulate_grad_batches=8,
    val_check_interval=0.25,
    progress_bar_refresh_rate=1,
    weights_summary='top',

when I check GPU utilization it is running correctly

Expected behaviour

a properly updated progress bar

Environment

Note: Bugs with code are solved faster ! Colab Notebook should be made public !

cc @tchaton @rohitgr7 @akihironitta @awaelchli

[tchaton]

Hey @Borda. I think this one is a P0 as it is a very common usage.

[Borda]

@awaelchli just testing with v1.3.1 and seems like the training does not start at all :/

[awaelchli]

I did a quick test on a Grid session with this snippet:

import torch.multiprocessing as mp


def foo(process_idx):
    print(process_idx)


if __name__ == "__main__":
    mp.spawn(foo, nprocs=2)  # should print 0 1

It crashes in Jupyter environment and works fine as a regular python program.

---------------------------------------------------------------------------
ProcessExitedException                    Traceback (most recent call last)
<ipython-input-20-53d29309dc61> in <module>
      7 
      8 if __name__ == "__main__":
----> 9     mp.spawn(foo, nprocs=2)  # should print 0 1

~/conda/lib/python3.7/site-packages/torch/multiprocessing/spawn.py in spawn(fn, args, nprocs, join, daemon, start_method)
    228                ' torch.multiprocessing.start_process(...)' % start_method)
    229         warnings.warn(msg)
--> 230     return start_processes(fn, args, nprocs, join, daemon, start_method='spawn')

~/conda/lib/python3.7/site-packages/torch/multiprocessing/spawn.py in start_processes(fn, args, nprocs, join, daemon, start_method)
    186 
    187     # Loop on join until it returns True or raises an exception.
--> 188     while not context.join():
    189         pass
    190 

~/conda/lib/python3.7/site-packages/torch/multiprocessing/spawn.py in join(self, timeout)
    142                     error_index=error_index,
    143                     error_pid=failed_process.pid,
--> 144                     exit_code=exitcode
    145                 )
    146 

ProcessExitedException: process 1 terminated with exit code 1

So if this does not work in notebooks, then Lightning will also fail like this. Have you seen the same error?