Patches for Lightning have not kept up with backwards-incompatible changes

a-gardner1 commented 2 months ago

Describe the bug

While there has clearly been some effort to keep pace with changes to Lightning (see #1033), it has fallen behind since the initial patches were created (https://github.com/allegroai/clearml/commit/64e10b2f62f59244750e73604836b57470a2f0d7) and new versions of Lightning were released. Unfortunately, it silently fails to apply patches to model saving and restoration, which can hide the fact that model logging doesn't fully work as expected. One of the two related (and nearly duplicate) patch methods is shown below (linked here)

    @staticmethod
    def _patch_pytorch_lightning_io():
        if PatchPyTorchModelIO.__patched_pytorch_lightning:
            return

        if 'pytorch_lightning' not in sys.modules:
            return

        PatchPyTorchModelIO.__patched_pytorch_lightning = True

        # noinspection PyBroadException
        try:
            import pytorch_lightning  # noqa

            pytorch_lightning.trainer.Trainer.save_checkpoint = _patched_call(
                pytorch_lightning.trainer.Trainer.save_checkpoint, PatchPyTorchModelIO._save)  # noqa

            pytorch_lightning.trainer.Trainer.restore = _patched_call(
                pytorch_lightning.trainer.Trainer.restore, PatchPyTorchModelIO._load_from_obj)  # noqa
        except ImportError:
            pass
        except Exception:
            pass

        # noinspection PyBroadException
        try:
            import pytorch_lightning  # noqa

            # noinspection PyUnresolvedReferences
            pytorch_lightning.trainer.connectors.checkpoint_connector.CheckpointConnector.save_checkpoint = \
                _patched_call(
                    pytorch_lightning.trainer.connectors.checkpoint_connector.CheckpointConnector.save_checkpoint,
                    PatchPyTorchModelIO._save)  # noqa

            # noinspection PyUnresolvedReferences
            pytorch_lightning.trainer.connectors.checkpoint_connector.CheckpointConnector.restore = \
                _patched_call(
                    pytorch_lightning.trainer.connectors.checkpoint_connector.CheckpointConnector.restore,
                    PatchPyTorchModelIO._load_from_obj)  # noqa
        except ImportError:
            pass
        except Exception:
            pass

Three AttributeErrors exist in _patch_pytorch_lightning_io with newer versions of pytorch-lightning:

In pytorch-lightning-0.10.0, Trainer.restore was removed when CheckpointConnector was introduced and the restore method was no longer inherited from TrainerIOMixin (https://github.com/Lightning-AI/pytorch-lightning/commit/4724cdf5e0dc938ebff0a6d6b4477eec99326542)
In pytorch-lightning-2.0.0, CheckpointConnector was renamed to _CheckpointConnector (https://github.com/Lightning-AI/pytorch-lightning/pull/17008)
In pytorch-lightning-2.1.0, _CheckpointConnector.save_checkpoint was removed and inlined into Trainer (https://github.com/Lightning-AI/pytorch-lightning/pull/17408#discussion_r1170415577)

To reproduce

No reproduction is necessary. There are multiple clear AttributeErrors that get caught by the Exception handler depending on the pytorch-lightning version.

Expected behaviour

The checkpointing mechanism of pytorch-lightning should have been patched to enable automatic logging of models with ClearML.

Environment

Server type: self hosted
ClearML SDK Version: 1.15.1
ClearML Server Version (Only for self hosted): WebApp: 1.11.0-000 • Server: 1.12.0- • API: 2.26
Python Version: 3.11
OS: Linux

ainoam commented 2 months ago

Thanks for letting us know @a-gardner1. We'll take a look and update on fix availability.

a-gardner1 commented 2 months ago

I can open a PR with a proposed fix if you like. I've already implemented one

ainoam commented 2 months ago

Contributions are most welcome @a-gardner1 🙂

a-gardner1 commented 1 month ago

@ainoam In case you missed it, I did open a PR. Let me know if anything looks off

ainoam commented 1 month ago

Thanks for the friendly nudge @a-gardner1. We'll try to address your PR soon.

allegroai / clearml