Repository: Lightning-AI/pytorch-lightning
Pull request: Raise exception when `load_from_checkpoint` called from instance
## What does this PR do?

<!--
Please include a summary of the change and which issue is fixed.
Please also include relevant motivation and context.
List any dependencies that are required for this change.

If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

The following links the related issue to the PR (https://docs.github.com/en/free-pro-team@latest/github/managing-your-work-on-github/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword)
-->

Fixes https://github.com/Lightning-AI/lightning/issues/18169

<!-- Does your PR introduce any breaking changes? If yes, please list them. -->

<details>
  <summary><b>Before submitting</b></summary>

- Was this **discussed/agreed** via a GitHub issue? (not for typos and docs)
- [x ] Did you read the [contributor guideline](https://github.com/Lightning-AI/lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [x ] Did you make sure your **PR does only one thing**, instead of bundling different changes together?
- Did you make sure to **update the documentation** with your changes? (if necessary)
  - Updated docstring
- Did you write any **new necessary tests**? (not for typos and docs)
  - I'm not away of a way to test for the warning we emit.
- [x ] Did you verify new and **existing tests pass** locally with your changes?
  - Tests were broken on master on my machine prior to my changes.
- Did you list all the **breaking changes** introduced by this pull request?
- Did you **update the CHANGELOG**? (not for typos, docs, test updates, or minor internal changes/refactors)


<!-- In the CHANGELOG, separate each item in the unreleased section by a blank line to reduce collisions -->

</details>

## PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing, make sure you have read the [review guidelines](https://github.com/Lightning-AI/lightning/wiki/Review-guidelines). In short, see the following bullet-list:

<details>
  <summary>Reviewer checklist</summary>

- [ ] Is this pull request ready for review? (if not, please submit in draft mode)
- [ ] Check that all items from **Before submitting** are resolved
- [ ] Make sure the title is self-explanatory and the description concisely explains the PR
- [ ] Add labels and milestones (and optionally projects) to the PR so it can be classified

</details>

<!--

Did you have fun?

Make sure you had fun coding 🙃

-->


Problem statement:
`load_from_checkpoint` should not be called as an instance method
### Outline & Motivation

Spent a while debugging : 
```python
model.load_from_checkpoint(model_path)
```
which should be : 
```python
model = LitModel.load_from_checkpoint(model_path)
```
there should be a warning when the instance method is called. 

### Pitch

```python
model.load_from_checkpoint(model_path)
```
which should be : 
```python
model = LitModel.load_from_checkpoint(model_path)
```
there should be a warning when the first approach i.e. instance method is called. 

### Additional context

_No response_

cc @justusschock @awaelchli @carmocca @borda

Hints from issue discussion:
Thanks for the suggestion @ganessh22 
We already have an error message for this, but I think we could improve it to give some instructions what to do instead:
https://github.com/Lightning-AI/lightning/blob/df4ad1932b0ad9a5ffe9b335726e4d12f6820f23/src/lightning/pytorch/core/saving.py#L99

I see. Mistake from my side was assuming behaviour similar to Pytorch's `load_state_dict` which does inplace. If I do :
```python
model = model.load_from_checkpoint(model_path)
```
it works fine. I can open a PR if you suggest the improved instructions. I am familiar with the codebase. I think logging to not call inplace would be nice. 
@ganessh22 You are very welcome to open a PR. My suggestion would be to edit this error message here

https://github.com/Lightning-AI/lightning/blob/df4ad1932b0ad9a5ffe9b335726e4d12f6820f23/src/lightning/pytorch/core/saving.py#L99

to this:

```py
if isinstance(cls, LightningModule):
    raise ValueError(
        f"Calling `.load_from_checkpoint` on a model instance is not supported. Please change it to `{type(cls).__name__}.load_from_checkpoint(...)"
```





> @ganessh22 You are very welcome to open a PR. My suggestion would be to edit this error message here
> 
> https://github.com/Lightning-AI/lightning/blob/df4ad1932b0ad9a5ffe9b335726e4d12f6820f23/src/lightning/pytorch/core/saving.py#L99
> 
> to this:
> 
> ```python
> if isinstance(cls, LightningModule):
>     raise ValueError(
>         f"Calling `.load_from_checkpoint` on a model instance is not supported. Please change it to `{type(cls).__name__}.load_from_checkpoint(...)"
> ```

This method doesn't actually work, since `_load_from_checkpoint` is called inside the `classmethod` `load_from_checkpoint`, and the first argument (`cls`) is guaranteed to be the subbclass type and not the instance. 
I found a more generic way to solve the problem of guaranteeing a classmethod is only called on the class and not an instance: we define our own `classmethod` decorator. Let's called it `restricted_classmethod`:

```python
class restricted_classmethod:
    def __init__(self, method):
        self.method = method
    def __get__(self, instance, cls):
        if instance is not None:
            # print(f"Warning: classmethod {cls.__name__}.{self.method.__name__} called on an instance.")
            raise TypeError(f"restricted classmethod {cls.__name__}.{self.method.__name__} cannot be called from an instance. Please call it from the class type.")
        return lambda *args, **kwargs: self.method(cls, *args, **kwargs)


class A:
    @restricted_classmethod
    def cm(cls):
        print("Hello from inside classmethod")

a = A()
a.cm()  # Raises TypeError

A.cm()  # Works fine.
```

However, this limits the usage of the classmethod, and I'm not sure if this extra guard rail is worth.

PS: the best solution is if we can mark this method as [`[[nodiscard]]`](https://en.cppreference.com/w/cpp/language/attributes/nodiscard), but that's unfortunately not possible in Python. 
@kwsp This was an oversight on my part. You are right the suggestion I made won't work and now it also becomes clear to me why we haven't done this already in years 😅 

> However, this limits the usage of the classmethod

Could you explain the cons of the `restricted_classmethod` vs. Python's?

Given that this UX question came up several times already in the past, we should definitely consider replacing `classmethod` with such a restricted decorator.
If we raise an exception, the only con I can think of would be that it precludes this (rare) use case: 

```python
model2 = model.load_checkpoint(...)
```

To prevent breaking changes, we should probably emit a warning instead of raising an exception:

```python
import logging

log = logging.getLogger(__name__)

class restricted_classmethod:
    def __init__(self, method):
        self.method = method
    def __get__(self, instance, cls):
        if instance is not None:
            # Replace with the appropriate logger
            log.warning(f"classmethod {cls.__name__}.{self.method.__name__} called on an instance.")
        return lambda *args, **kwargs: self.method(cls, *args, **kwargs)

```

If we want to implement `restricted_classmethod`, let's first compile a list of classmethods in the lightning code base that would benefit from this. 
Adding my 2 cents here: I think the restricted_classmethod is great but it's an unnecessary abstraction. In this case I'd rather have an explicit type check followed by an eventual error in combination with a regular class method -- even at the cost of potential code duplication.

I think this construct would just hurt readability and debugging.
> In this case I'd rather have an explicit type check followed by an eventual error in combination with a regular class method

How would it look like? As pointed out earlier, the check can't be done in the `load_from_checkpoint` function directly as we never get access to the instance. If there is an easier way of doing it without introducing a context manager, I would also prefer that.
@justusschock I'd like to see your proposal. The reason I proposed the above is because in a `classmethod`, you don't have access to the instance even if the method is called on an instance. `classmethod`s and instance methods live in the same namespace, where the classmethod explicitly ignores the instance even if its present. The only way to have a method recognize whether its called on the type or an instance is by [implementing descriptors](https://docs.python.org/3/reference/datamodel.html#implementing-descriptors), as I proposed above.

Commit message: changed test for restricted_classmethod.

Patch to review:
diff --git a/src/lightning/pytorch/core/datamodule.py b/src/lightning/pytorch/core/datamodule.py
index bda6dac711d06..fa081b6f9d8f3 100644
--- a/src/lightning/pytorch/core/datamodule.py
+++ b/src/lightning/pytorch/core/datamodule.py
@@ -24,6 +24,7 @@
 from lightning.pytorch.core.hooks import DataHooks
 from lightning.pytorch.core.mixins import HyperparametersMixin
 from lightning.pytorch.core.saving import _load_from_checkpoint
+from lightning.pytorch.utilities.restricted_classmethod import restricted_classmethod
 from lightning.pytorch.utilities.types import EVAL_DATALOADERS, TRAIN_DATALOADERS
 
 
@@ -157,7 +158,7 @@ def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
         """
         pass
 
-    @classmethod
+    @restricted_classmethod
     def load_from_checkpoint(
         cls,
         checkpoint_path: Union[_PATH, IO],
@@ -200,8 +201,9 @@ def load_from_checkpoint(
             :class:`LightningDataModule` instance with loaded weights and hyperparameters (if available).
 
         Note:
-            ``load_from_checkpoint`` is a **class** method. You should use your :class:`LightningDataModule`
-            **class** to call it instead of the :class:`LightningDataModule` instance.
+            ``load_from_checkpoint`` is a **class** method. You must use your :class:`LightningDataModule`
+            **class** to call it instead of the :class:`LightningDataModule` instance, or a
+            ``TypeError`` will be raised.
 
         Example::
 
@@ -223,7 +225,7 @@ def load_from_checkpoint(
 
         """
         loaded = _load_from_checkpoint(
-            cls,
+            cls,  # type: ignore
             checkpoint_path,
             map_location=map_location,
             hparams_file=hparams_file,
diff --git a/src/lightning/pytorch/core/module.py b/src/lightning/pytorch/core/module.py
index 00f2f6bec401c..73daadf93da92 100644
--- a/src/lightning/pytorch/core/module.py
+++ b/src/lightning/pytorch/core/module.py
@@ -65,6 +65,7 @@
 from lightning.pytorch.utilities.exceptions import MisconfigurationException
 from lightning.pytorch.utilities.imports import _TORCHMETRICS_GREATER_EQUAL_0_9_1
 from lightning.pytorch.utilities.rank_zero import rank_zero_debug, rank_zero_warn, WarningCache
+from lightning.pytorch.utilities.restricted_classmethod import restricted_classmethod
 from lightning.pytorch.utilities.signature_utils import is_param_in_hook_signature
 from lightning.pytorch.utilities.types import _METRIC, LRSchedulerPLType, LRSchedulerTypeUnion, STEP_OUTPUT
 
@@ -1459,7 +1460,7 @@ def forward(self, x):
 
         return torchscript_module
 
-    @classmethod
+    @restricted_classmethod
     def load_from_checkpoint(
         cls,
         checkpoint_path: Union[_PATH, IO],
@@ -1507,7 +1508,8 @@ def load_from_checkpoint(
 
         Note:
             ``load_from_checkpoint`` is a **class** method. You should use your :class:`LightningModule`
-            **class** to call it instead of the :class:`LightningModule` instance.
+            **class** to call it instead of the :class:`LightningModule` instance, or a
+            ``TypeError`` will be raised.
 
         Example::
 
@@ -1540,7 +1542,7 @@ def load_from_checkpoint(
             y_hat = pretrained_model(x)
         """
         loaded = _load_from_checkpoint(
-            cls,
+            cls,  # type: ignore
             checkpoint_path,
             map_location,
             hparams_file,
diff --git a/src/lightning/pytorch/utilities/restricted_classmethod.py b/src/lightning/pytorch/utilities/restricted_classmethod.py
new file mode 100644
index 0000000000000..9e51e2587f778
--- /dev/null
+++ b/src/lightning/pytorch/utilities/restricted_classmethod.py
@@ -0,0 +1,23 @@
+from typing import Callable, Concatenate, Generic, ParamSpec, TypeVar
+
+_T = TypeVar("_T")  # type of the method owner
+_P = ParamSpec("_P")  # parameters of the decorated method
+_R_co = TypeVar("_R_co", covariant=True)  # return type of the decorated method
+
+
+class restricted_classmethod(Generic[_T, _P, _R_co]):
+    """
+    Custom `classmethod` that emits a warning when the classmethod is
+    called on an instance and not the class type.
+    """
+
+    def __init__(self, method: Callable[Concatenate[_T, _P], _R_co]) -> None:
+        self.method = method
+
+    def __get__(self, instance: _T, cls: type[_T]) -> Callable[_P, _R_co]:
+        if instance is not None:
+            raise TypeError(
+                f"The classmethod `{cls.__name__}.{self.method.__name__}` cannot be invoked on an instance. Please "
+                f"invoke it on the class type and make sure the return value is used."
+            )
+        return self.method.__get__(cls, cls)
diff --git a/tests/tests_pytorch/utilities/test_restricted_classmethod.py b/tests/tests_pytorch/utilities/test_restricted_classmethod.py
new file mode 100644
index 0000000000000..1606c47726230
--- /dev/null
+++ b/tests/tests_pytorch/utilities/test_restricted_classmethod.py
@@ -0,0 +1,18 @@
+import pytest
+
+from lightning.pytorch.utilities.restricted_classmethod import restricted_classmethod
+
+
+class _C:
+    @restricted_classmethod
+    def _rcmethod(cls):
+        pass
+
+
+def test_restricted_classmethod():
+    c = _C()
+
+    with pytest.raises(TypeError):
+        c._rcmethod()
+
+    _C._rcmethod()  # no exception


Code review comments:
diff_hunk: @@ -0,0 +1,18 @@
+import pytest
+
+from lightning.pytorch.utilities.restricted_classmethod import restricted_classmethod
+
+
+class _C:
path: tests/tests_pytorch/utilities/test_restricted_classmethod.py
text: @user1:
Let's use meaningful names, for example
```suggestion
class Restricted:
```

diff_hunk: @@ -0,0 +1,18 @@
+import pytest
+
+from lightning.pytorch.utilities.restricted_classmethod import restricted_classmethod
+
+
+class _C:
+    @restricted_classmethod
+    def _rcmethod(cls):
+        pass
+
+
+def test_restricted_classmethod():
+    c = _C()
path: tests/tests_pytorch/utilities/test_restricted_classmethod.py
text: @user1:
```suggestion
    restricted = Restricted()
```

diff_hunk: @@ -0,0 +1,18 @@
+import pytest
+
+from lightning.pytorch.utilities.restricted_classmethod import restricted_classmethod
+
+
+class _C:
+    @restricted_classmethod
+    def _rcmethod(cls):
+        pass
+
+
+def test_restricted_classmethod():
+    c = _C()
+
+    with pytest.raises(TypeError):
path: tests/tests_pytorch/utilities/test_restricted_classmethod.py
text: @user1:
```suggestion
    with pytest.raises(TypeError, match="cannot be called on an instance"):
```

path: src/lightning/pytorch/utilities/restricted_classmethod.py
text: @user1:
Perhaps we don't need to put this in a new file. Would it fit into utilities/model_helpers.py?

diff_hunk: @@ -0,0 +1,18 @@
+import pytest
+
+from lightning.pytorch.utilities.restricted_classmethod import restricted_classmethod
+
+
+class _C:
+    @restricted_classmethod
+    def _rcmethod(cls):
path: tests/tests_pytorch/utilities/test_restricted_classmethod.py
text: @user1:
```suggestion
    def restricted_method(cls):
```

diff_hunk: @@ -0,0 +1,23 @@
+from typing import Callable, Concatenate, Generic, ParamSpec, TypeVar
+
+_T = TypeVar("_T")  # type of the method owner
+_P = ParamSpec("_P")  # parameters of the decorated method
+_R_co = TypeVar("_R_co", covariant=True)  # return type of the decorated method
+
+
+class restricted_classmethod(Generic[_T, _P, _R_co]):
path: src/lightning/pytorch/utilities/restricted_classmethod.py
text: @user1:
I suggest to make this protected so we can change it in the future if needed. 

```suggestion
class _restricted_classmethod(Generic[_T, _P, _R_co]):
```

diff_hunk: @@ -0,0 +1,23 @@
+from typing import Callable, Concatenate, Generic, ParamSpec, TypeVar
+
+_T = TypeVar("_T")  # type of the method owner
+_P = ParamSpec("_P")  # parameters of the decorated method
+_R_co = TypeVar("_R_co", covariant=True)  # return type of the decorated method
+
+
+class restricted_classmethod(Generic[_T, _P, _R_co]):
+    """
+    Custom `classmethod` that emits a warning when the classmethod is
+    called on an instance and not the class type.
+    """
+
+    def __init__(self, method: Callable[Concatenate[_T, _P], _R_co]) -> None:
+        self.method = method
+
+    def __get__(self, instance: _T, cls: type[_T]) -> Callable[_P, _R_co]:
+        if instance is not None:
+            raise TypeError(
+                f"The classmethod `{cls.__name__}.{self.method.__name__}` cannot be invoked on an instance. Please "
+                f"invoke it on the class type and make sure the return value is used."
path: src/lightning/pytorch/utilities/restricted_classmethod.py
text: @user1:
I propose simpler language
```suggestion
                f"The classmethod `{cls.__name__}.{self.method.__name__}` cannot be called on an instance. Please "
                f"call it on the class type and make sure the return value is used."
```

diff_hunk: @@ -0,0 +1,23 @@
+from typing import Callable, Concatenate, Generic, ParamSpec, TypeVar
+
+_T = TypeVar("_T")  # type of the method owner
+_P = ParamSpec("_P")  # parameters of the decorated method
+_R_co = TypeVar("_R_co", covariant=True)  # return type of the decorated method
+
+
+class restricted_classmethod(Generic[_T, _P, _R_co]):
+    """
+    Custom `classmethod` that emits a warning when the classmethod is
path: src/lightning/pytorch/utilities/restricted_classmethod.py
text: @user1:
```suggestion
    Custom `classmethod` that raises an exception when the classmethod is
```