Replace _DataModuleWrapper with __new__ [1/2]

[carmocca]

What does this PR do?

We now rely on __new__ to wrap and track the DataHooks calls.

Fixes:

This fixes a bug where the _DataModuleWrapper implementation was adding as many wrapped layers as the number of DataModule subclasses.
This is not a problem if a wrapped layer does not impact any further wrapper layers which is no longer the case after #7238 (blocked by this PR).
And even without #7238, we should not wrap these hooks multiple times.

To compare the behaviors, see the following minimal examples:

Current master

import functools
import pickle


class Metaclass(type):
    def __init__(cls, *args, **kwargs):
        super().__init__(*args, **kwargs)
        cls.has_wrapped = False

    def __call__(cls, *args, **kwargs):
        if not cls.has_wrapped:
            cls.thing = cls.wrapper(cls.thing)
            cls.has_wrapped = True
        obj = type.__call__(cls, *args, **kwargs)
        return obj

    @staticmethod
    def wrapper(fn):
        @functools.wraps(fn)
        def wrapped_fn(*args, **kwargs):
            print(f"Wrapped {fn.__name__}. args: {args}, kwargs: {kwargs}")
            return fn(*args, **kwargs)
        return wrapped_fn


class Foo(metaclass=Metaclass):
    def thing(self):
        print(f"{self.__class__.__name__}.thing()")

class Bar(Foo): ...
class FooBar(Bar): ...


f = Foo()
b = Bar()
fb = FooBar()

pickle.dumps(f)

f.thing()
b.thing()
fb.thing()
"""
Wrapped thing. args: (<__main__.Foo object at 0x10f6ef940>,), kwargs: {}
Foo.thing()
Wrapped thing. args: (<__main__.Bar object at 0x10f6ef820>,), kwargs: {}
Wrapped thing. args: (<__main__.Bar object at 0x10f6ef820>,), kwargs: {}
Bar.thing()
Wrapped thing. args: (<__main__.FooBar object at 0x10f6ef8e0>,), kwargs: {}
Wrapped thing. args: (<__main__.FooBar object at 0x10f6ef8e0>,), kwargs: {}
Wrapped thing. args: (<__main__.FooBar object at 0x10f6ef8e0>,), kwargs: {}
FooBar.thing()
"""

This PR

import functools
import pickle


class Foo:
    def __new__(cls, *args, **kwargs):
        obj = super(Foo, cls).__new__(cls)
        obj.__args = args
        obj.__kwargs = kwargs
        obj.thing = cls.wrapper(obj, obj.thing)
        return obj

    def __getstate__(self):
        # avoids _pickle.PicklingError: Can't pickle <...>: it's not the same object as <...>
        d = self.__dict__.copy()
        del d["thing"]
        return d

    @staticmethod
    def wrapper(obj, fn):
        @functools.wraps(fn)
        def wrapped_fn(*args, **kwargs):
            print(f"Wrapped {obj}.{fn.__name__}. args: {args}, kwargs: {kwargs}")
            return fn(*args, **kwargs)
        return wrapped_fn

    def thing(self):
        print(f"{self.__class__.__name__}.thing()")

class Bar(Foo): ...
class FooBar(Bar): ...


f = Foo()
b = Bar()
fb = FooBar()

# needs `__getstate__`
pickle.dumps(f)

f.thing()
b.thing()
fb.thing()
"""
Wrapped <__main__.Foo object at 0x1075bb970>.thing. args: (), kwargs: {}
Foo.thing()
Wrapped <__main__.Bar object at 0x1075bb910>.thing. args: (), kwargs: {}
Bar.thing()
Wrapped <__main__.FooBar object at 0x1074195b0>.thing. args: (), kwargs: {}
FooBar.thing()
"""

Question:

How do I even test this automatically?

Before submitting

• Was this discussed/approved via a GitHub issue? (not for typos and docs)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• [n/a] 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)
• Did you verify new and existing tests pass locally with your changes?
• [n/a] Did you update the CHANGELOG? (not for typos, docs, test updates, or internal minor changes/refactorings)

PR review

• 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

[codecov]

Codecov Report

Merging #7289 (905a0bf) into master (338f5a3) will decrease coverage by 0%.
The diff coverage is 100%.

@@          Coverage Diff           @@
##           master   #7289   +/-   ##
======================================
- Coverage      91%     91%   -0%     
======================================
  Files         199     199           
  Lines       12797   12807   +10     
======================================
- Hits        11678   11672    -6     
- Misses       1119    1135   +16     

[ananthsub]

Would it be simpler to have explicit properties and setters on the datamodule which the trainer can set after calling the hooks?

I also am not sure if the current properties are correct for setup and teardown if one re-uses the same module across consecutive calls.

What's expected for this?

trainer.fit(model, dm)
trainer.fit(model, dm)

Should the second call skip calling the datamodule hooks?

[carmocca]

Would it be simpler to have explicit properties and setters on the datamodule which the trainer can set after calling the hooks?

This is the case already. But instead of setting them manually by the trainer, it is done automatically using these "hacks".

Should the second call skip calling the datamodule hooks?

This is also the case already. It is not if the user does the call manually (e.g. dm.prepare_data()), and for that, I have #7238

[ananthsub]

Actually why do we need these setters at all? Couldn't we push the responsibility to the datamodule on whether it needs to download data again? Then we can simplify the trainer and avoid all the hook wrappers here

This way we have to worry less about state across trainer calls

[carmocca]

If I had designed this from the beginning I would have done one of the 2 following options instead of what we have:

1. Move the responsability to the user so they check if the files are already downloaded and splits already created (if they want)
2. Ask users to call super().setup(stage) and in the base class we set these has_setup_stage checks.

But this is not the current design and I don't think we can safely switch the behaviour.

[justusschock]

I left some comments. It looks good to me. However, I don't enough about these things for being confident enough to approve.

[ananthsub]

But this is not the current design and I don't think we can safely switch the behaviour.

Do you think we can add warnings in the trainer if the has_{setup/teardown}_* properties are True? And then we can phase out these checks in v1.5?

[carmocca]

But this is not the current design and I don't think we can safely switch the behaviour.

Do you think we can add warnings in the trainer if the has_{setup/teardown}_* properties are True? And then we can phase out these checks in v1.5?

Do you mean if we go with option (1)?:

Move the responsability to the user so they check if the files are already downloaded and splits already created (if they want)

Do you want to open an RFC or discussion about this?