Crash on multiple initializations in load_checkpoint() calls

[pedrocolon93]

🐛 Bug

I am getting an exception in:
migration.py

Particularly with the line:
del sys.modules["pytorch_lightning.utilities.argparse_utils"]

A cheap patch was to try:/catch the line.

To Reproduce

I am using a threadexecutor to init 4 threads to run a model on each thread (each with a cuda device). At some point between thread loads of checkpoints i get a KeyError in the line I mentioned above. I think that on the first load it deletes this module, but doesnt check if its deleted in subsequent loads.

Expected behavior

That the checkpoints load in the multiple threads.

Environment

• CUDA:
  • GPU:
    • NVIDIA RTX A6000
    • NVIDIA RTX A6000
    • NVIDIA RTX A6000
    • NVIDIA RTX A6000
  • available: True
  • version: 11.3
• Packages:
  • numpy: 1.21.2
  • pyTorch_debug: False
  • pyTorch_version: 1.10.0
  • pytorch-lightning: 1.6.0
  • tqdm: 4.62.3
• System:
  • OS: Linux
  • architecture:
    • 64bit
    • ELF
  • processor: x86_64
  • python: 3.9.7
  • version: Names of parameters may benefit from not being abbreviated #119-Ubuntu SMP Mon Mar 7 18:49:24 UTC 2022

Additional context

The idea was to have each model in a thread feeding it data.

cc @awaelchli @rohitgr7 @akihironitta @ananthsub @ninginthecloud

[krshrimali]

Hi, @pedrocolon93 - Thank you for creating this issue. Just to let you know, I'm taking a look at it and trying to reproduce it first on my system. :) If you have any code sample to share that will help reproduce, it would be helpful, but if not - that's also okay. I'll keep you updated here. ⚡

[krshrimali]

Update: can reproduce this locally. Will create a fix today! :)

[krshrimali]

Hi, @pedrocolon93 - Upon further investigation, I believe it's more about fixing race conditions with threading. You can use and acquire a lock and release once you are done loading the checkpoint in your function. Example:

import threading
from pytorch_lightning.utilities.migration import pl_legacy_patch
import torch

def model_load(lock):
    lock.acquire()
    with pl_legacy_patch():
        m = torch.load("checkpoint_version_1.1.6.ckpt")
    lock.release()
    return m

def main():
    lock = threading.Lock()
    t1 = threading.Thread(target=model_load, args=(lock,))
    t2 = threading.Thread(target=model_load, args=(lock,))

    t1.start()
    t2.start()

    t1.join()
    t2.join()

Or you can use a ThreadPool and use apply_async method. This should hopefully solve the issues with you. Let me know if it doesn't. :)

[carmocca]

I am using a threadexecutor to init 4 threads to run a model on each thread (each with a cuda device).

@pedrocolon93 Could you share more details and ideally some code about your use case? Manually handling threads is unusual and not recommended by the framework. Is there any reason why not enable Trainer(strategy="ddp_spawn") to let the framework do this for you?

Depending on your code structure, the above suggestion by @krshrimali is good.

[pedrocolon93]

Hi! I was loading a model checkpoint for inference (loading a language model for language generation), so I was not even using the trainer, but the above suggestion seems to work. Seems like an odd issue to have either way; I thought model loading was independent/could be handled with multi-threading.

[carmocca]

I think that's a usecase that should be supported then. I'm thinking that this logic should only be done on rank zero since sys.modules are shared

cc @awaelchli

[awaelchli]

I didn't consider this case, but it's true that modifying the global is not thread safe I guess. What we could do, if I understand correctly, is add the thread lock into the pl_legacy_patch context manager.