Checkpoint loading patch fails when using a pre-trained featurizer within a Lightning module.

[jdhorwood]

🐛 Bug

I am getting the error: ModuleNotFoundError: No module named 'pytorch_lightning.utilities.argparse_utils' coming from
migration.py when loading checkpoints for models which leverage pre-trained models also trained via Pytorch Lightning.

Essentially, the context manager is called recursively, since the model relies on features generated from a pre-trained model. This model also requires a call to load_from_checkpoint(). As a result, the exit function of the context manager gets called twice in a row, with the error occurring on the second call.

Expected behavior

I would expect checkpoint loading to not raise en error because of a pre-trained featurizer model.

Environment

pytorch-lightning == 1.6.1

cc @awaelchli @ananthsub @ninginthecloud @rohitgr7

[jdhorwood]

A solution to this in the legacy_patch would be to simply check if pytorch_lightning.utilities.argparse_utils is still in sys.modules before the del command.

[carmocca]

@JDavid04 Feel free to open a PR with the proposed fix

[jdhorwood]

I am happy to make a PR with the fix. I think I might need permission to push in order to create the PR.

I've tested the fix locally off the recent main branch and it works, but I encountered new issues arising from this PR when running in an already parallel environment (in my case, with Ray Tune). In that case, the code never terminates presumably because of the global thread lock.

So I would propose removing the thread locking and keeping only the check. I tried to test locally that the threading test passes with this fix also but the tests also seem to have issues with the paths not being correct. I am not sure if that's related to this PR.

[rohitgr7]

I am happy to make a PR with the fix. I think I might need permission to push in order to create the PR

you can fork the repo and send the fix?

[carmocca]

Hi @JDavid04. Is this still an issue? Are you able to finish your PR?