Repository: Lightning-AI/pytorch-lightning
Pull request: Add `is_wrapped` utility function for Fabric
## What does this PR do?

Method to check whether a Fabric model was set up already #16484 


Fixes #16484 

<!-- 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)
- Did you write any **new necessary tests**? (not for typos and docs)
- [x] Did you verify new and **existing tests pass** locally with your 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>

- [x] Is this pull request ready for review? (if not, please submit in draft mode)
- [x] Check that all items from **Before submitting** are resolved
- [x] Make sure the title is self-explanatory and the description concisely explains the PR
- [x] 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:
Method to check whether a Fabric model was set up already
### Description & Motivation

Currently, if you pass a model to Fabric.setup() that has already been wrapped/set up, you get an error. This is currently expected, because there are several reasons and edge cases where this would be invalid usage. However, the only way to check whether a model is set up already is with an instance check on a protected class:

```py
already_set_up = isinstance(model, _FabricModel)
```

There are valid use case where such a check comes up. For example:

You build a trainer that has a fit method with the following implementation.

```py
model, optimizer = fabric.setup(model, optimizer)
for epoch in range(epochs):
    self.train(model, optimizer)
    self.validate(model)
```

This trainer has a validate() method that can either be called separately, or internally within the fit loop. Hence, the validate() method needs to check whether the model was already set up.




### Pitch

Introduce `Fabric.is_setup()`, or function `is_setup()`.

```py
assert not is_setup(model)
model = fabric.setup(model)
assert is_setup(model)
```

### Alternatives

There are two alternatives:

**Alternative 1**
Make the `Fabric.setup()` a no-op when a model is passed in that was already set up. However, there are problematic edge cases, that break the contract with the strategies. For example, when setup is called with a model that is already set up, but an optimizer that is not set up yet:
```py
model = setup(model)
model = setup(model, optimizer)  # what should happen?
```
OR
If a model is set up twice but with different optimizers
```py
model, optimizer1 = setup(model, optimizer1)
model, optimizer2 = setup(model, optimizer2)
```

**Alternative 2**
The example I gave with the shared validation loop could also be re-written like this:

```py
def train(model, optimizer):
    ...

def validate(model):
    model = fabric.setup(model)
    return self._validate(model)
    
def _validate(model):
    # this assumes the passed in model is already set up

def fit(model):
    model, optimizer = fabric.setup(model, optimizer)
    for epoch in range(epochs):
        self.train(model, optimizer)
        self._validate(model)  # call _validate(), not validate()

```

### Additional context

cc @borda @tchaton @justusschock @awaelchli @carmocca @pietrolesci

Hints from issue discussion:
Would a property work too?

Do we need separate `is_*_setup` functions for the different components that can be set-up?
A property wouldn't work. You wouldn't be able to check `model.is_setup()` if model is an nn.Module (and therefore not set up). It would either have to be a function

```py
def is_setup(module) -> bool:
    return isinstance(module, _FabricModule)
```
or a (static) method on the Fabric object:
```py
def is_setup(self, module) -> bool:
    return isinstance(module, _FabricModule)
```



Is this open can I take this up ? Also amazing work with fabric 📦 
Thank you @Atharva-Phatak! Happy to let you take this. We need to decide whether to make this a function or a method. What are your thoughts? Also @justusschock @carmocca I would make it a function.

On another note, I'm generally not so pleased about the name of "setup". We use it as an action/verb, but it is backwards because in English, the verb is actually spelled "set up", and "setup" is the noun. This makes it awkward when using Fabric.setup, but is even more cringe when we have something like `is_setup` as a function name 😅 
A function would be much better since it would be independent of the fabric class. For the name, may I suggest the name ```check_state``` ? 
I agree with choosing a function. I wouldn't go with changing the "setup" name. It's widely used in software projects, including this one. It's familiar even if it's not the most natural English name to read
I also agree with this being a function. Similar to @carmocca , I'd keep the `setup` terminology.

@Atharva-Phatak with state usually we mean something different as this isn't really about the internal state but more whether a specific function call has already happened. 
Also it should be something with `is_X` so that returning a bool actually makes sense. 
Hi there, my two cents: `lightning.fabric.is_wrapped`.

I think that a module-level function (like in `torch.is_tensor`) would be great and I think that `is_wrapped` is a pretty explanatory name for it since we are wrapping `nn.Module`s in Fabric (also, "wrapping in fabric" has a nice physical sound to it :)).

```python
def is_wrapped(module: Union[nn.Module, _FabricModule]) -> bool:
    ...
```

  
I think this is a good alternative too. I like it.
@awaelchli Sounds good. I will get cracking on this.

Commit message: [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Patch to review:
diff --git a/src/lightning/fabric/__init__.py b/src/lightning/fabric/__init__.py
index c7084388528a8..972be1e488ee4 100644
--- a/src/lightning/fabric/__init__.py
+++ b/src/lightning/fabric/__init__.py
@@ -26,7 +26,7 @@
 os.environ["PYTORCH_NVML_BASED_CUDA_CHECK"] = "1"
 
 
-from lightning.fabric.fabric import Fabric  # noqa: E402
+from lightning.fabric.fabric import Fabric, is_wrapped  # noqa: E402
 from lightning.fabric.utilities.seed import seed_everything  # noqa: E402
 
 __all__ = ["Fabric", "seed_everything"]
diff --git a/src/lightning/fabric/fabric.py b/src/lightning/fabric/fabric.py
index c6b9423b400f2..c34e8b30dfdac 100644
--- a/src/lightning/fabric/fabric.py
+++ b/src/lightning/fabric/fabric.py
@@ -799,5 +799,13 @@ def _is_using_cli() -> bool:
     return bool(int(os.environ.get("LT_CLI_USED", "0")))
 
 
+def is_wrapped(module: Union[nn.Module, _FabricModule]) -> bool:
+    """Method to check if nn.Module is already wrapped as a FabricModule
+    Args:
+        module: The module to be checked if it is _FabricModule or not
+    """
+    return isinstance(module, _FabricModule)
+
+
 def _do_nothing(*_: Any) -> None:
     pass


Code review comments:
diff_hunk: @@ -799,5 +799,13 @@ def _is_using_cli() -> bool:
     return bool(int(os.environ.get("LT_CLI_USED", "0")))
 
 
+def is_wrapped(module: Union[nn.Module, _FabricModule]) -> bool:
path: src/lightning/fabric/fabric.py
text: @user1:
I know this is a very simple function, but could we still get a test for it?

@user2:
+1

I would also move the utility function to wrappers.py instead of fabric.py :) 

@author:
Sounds good. Will move the function there :)

diff_hunk: @@ -799,5 +799,13 @@ def _is_using_cli() -> bool:
     return bool(int(os.environ.get("LT_CLI_USED", "0")))
 
 
+def is_wrapped(module: Union[nn.Module, _FabricModule]) -> bool:
+    """Method to check if nn.Module is already wrapped as a FabricModule
+    Args:
+        module: The module to be checked if it is _FabricModule or not
path: src/lightning/fabric/fabric.py
text: @user1:
Can we extend this to an arbitrary object and not only check for _FabricModule, but also for _FabricOptimizer etc.?
That would be kind of consistent with `_unwrap_objects`

@user2:
+1

It could check for _FabricModule, _FabricOptimizer, and _FabricDataLoader

@author:
@user2, would you like? function for each check ?? 

@user2:
No, I think just extending the check inside this function is the best :))