KeyError: 'self' in save_hyperparameters() when custom metaclass used

[aditya0by0]

Bug description

I'm working with LightningDataModule and wanted to ensure that a method (_after_init) runs only once after full initialization, regardless of subclassing. For that, I implemented a custom metaclass (_InitMeta) that overrides __call__ to invoke _after_init after the instance is fully created.

When using create an instance of final subclass, I encounter a KeyError: 'self' inside save_hyperparameters().

What version are you seeing the problem on?

v2.5, v2.1

How to reproduce the bug

from typing import Any

from lightning import LightningDataModule


class _InitMeta(type):
    def __call__(
        cls: Any, *args: Any, **kwargs: Any
    ) -> Any:
        instance = super().__call__(*args, **kwargs)  # Create the instance
        if hasattr(instance, "_after_init"):
            instance._after_init(**kwargs)  # Call the method if defined
        return instance


class A(LightningDataModule, metaclass=_InitMeta):
    def __init__(self, *args, **kwargs):
        self.save_hyperparameters()
        self.a = 1
        self.b = 2
        super().__init__(*args, **kwargs)

    def print_ab(self, **kwargs: Any):
        print("in print ab")
        if kwargs.get("flag", False):
            print("flag is set to False")
            print("some other logic")
        else:
            print(self.a, self.b)

    def _after_init(self, **kwargs):
        """Called only once after full initialization."""
        self.print_ab(**kwargs)


class B(A):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.a += 1
        self.b += 2


class C(B):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.a += 1
        self.b += 2


if __name__ == "__main__":
    print("Creating C instance:")
    c = C()  # Should print 3, 6 only once

    print("\nCreating B instance:")
    b = B()  # Should print 2, 4 only once

    print("\nCreating A instance:")
    a = A()  # Should print 1, 2 only once

Error messages and logs

Creating C instance:
Traceback (most recent call last):
  File "G:\github-aditya0by0\python-chebai\test.py", line 48, in <module>
    c = C()  # Should print 3, 6 only once
  File "G:\github-aditya0by0\python-chebai\test.py", line 10, in __call__
    instance = super().__call__(*args, **kwargs)  # Create the instance
  File "G:\github-aditya0by0\python-chebai\test.py", line 41, in __init__
    super().__init__(**kwargs)
  File "G:\github-aditya0by0\python-chebai\test.py", line 34, in __init__
    super().__init__(**kwargs)
  File "G:\github-aditya0by0\python-chebai\test.py", line 18, in __init__
    self.save_hyperparameters()
  File "G:\anaconda3\envs\env_chebai\lib\site-packages\lightning\pytorch\core\mixins\hparams_mixin.py", line 112, in save_hyperparameters
    save_hyperparameters(self, *args, ignore=ignore, frame=frame)
  File "G:\anaconda3\envs\env_chebai\lib\site-packages\lightning\pytorch\utilities\parsing.py", line 165, in save_hyperparameters
    for local_args in collect_init_args(frame, [], classes=(HyperparametersMixin,)):
  File "G:\anaconda3\envs\env_chebai\lib\site-packages\lightning\pytorch\utilities\parsing.py", line 135, in collect_init_args
    return collect_init_args(frame.f_back, path_args, inside=True, classes=classes)
  File "G:\anaconda3\envs\env_chebai\lib\site-packages\lightning\pytorch\utilities\parsing.py", line 135, in collect_init_args
    return collect_init_args(frame.f_back, path_args, inside=True, classes=classes)
  File "G:\anaconda3\envs\env_chebai\lib\site-packages\lightning\pytorch\utilities\parsing.py", line 135, in collect_init_args
    return collect_init_args(frame.f_back, path_args, inside=True, classes=classes)
  File "G:\anaconda3\envs\env_chebai\lib\site-packages\lightning\pytorch\utilities\parsing.py", line 131, in collect_init_args
    local_self, local_args = _get_init_args(frame)
  File "G:\anaconda3\envs\env_chebai\lib\site-packages\lightning\pytorch\utilities\parsing.py", line 97, in _get_init_args
    local_args = {k: local_vars[k] for k in init_parameters}
  File "G:\anaconda3\envs\env_chebai\lib\site-packages\lightning\pytorch\utilities\parsing.py", line 97, in <dictcomp>
    local_args = {k: local_vars[k] for k in init_parameters}
KeyError: 'self'

Environment

Current environment

- PyTorch Lightning: 2.1.2
- Python: 3.10.14
- Torch: 2.5.1

More info

This is acknowledged as bug by @jsbueno on https://stackoverflow.com/questions/79554986/keyerror-self-in-save-hyperparameters-when-custom-metaclass-used-pytorch

• GitHub Discussion: KeyError: 'self' in save_hyperparameters() when custom metaclass used #20693

[jsbueno]

I have verified the error also occurs in current main head.

The introspection code used has two problems: it checks for parameters on a hardcoded __init__ named method (which makes sense), but them use the parameters name in the __init__ method to fetch local variable values in the method calling save_hyperparameters - if it was called from any other method than __init__, this error will raise.

And then, the second problem is that the backward tracking code to find all parameters used in the base classes doesn't distinguish the base class itself, and keeps back-tracking to frames before the first custom __init__ method. In this example, that is the metaclass __call__, and in the associated discussion, the method which created an instance of the class which called save_hyperparameters to be used as an associated attribute. (A scenario far more likely to happen).

[ved1beta]

The original issue occurred due to the interaction between PyTorch Lightning's save_hyperparameters() 
and custom metaclasses. Here's what was happening:

1. Problem:
   - save_hyperparameters() uses frame inspection to find initialization parameters
   - It walks up the call stack looking for __init__ methods
   - When a custom metaclass is used, the call stack includes the metaclass's __call__ method
   - This causes save_hyperparameters() to try to find 'self' in the metaclass's frame
   - Since metaclass __call__ is a class method, it doesn't have 'self' in its locals
   - This results in KeyError: 'self'

2. Solution:
   - Instead of relying on frame inspection, we store hyperparameters explicitly
   - The metaclass captures kwargs before instance creation
   - We don't pass kwargs through __init__ chain to avoid conflicts with LightningDataModule
   - After instance creation, we use _set_hparams() to properly set hyperparameters
   - We ensure _after_init is called exactly once using a flag

3. Benefits:
   - Avoids frame inspection issues
   - Maintains proper initialization order
   - Correctly handles hyperparameters
   - Prevents multiple initialization
   - Works with inheritance
"""

from typing import Any
from lightning import LightningDataModule

class _InitMeta(type):
    def __call__(cls: Any, *args: Any, **kwargs: Any) -> Any:
        # Create the instance first
        hparams = kwargs.copy()  # Save hyperparameters
        instance = super().__call__(*args)  # Don't pass kwargs to __init__
        
        # Store hyperparameters for later use
        instance._hparams_dict = hparams
        
        # Call _after_init if it exists, but only after the instance is fully initialized
        if hasattr(instance, "_after_init"):
            # Use a flag to prevent multiple calls
            if not getattr(instance, "_after_init_called", False):
                instance._after_init()
                instance._after_init_called = True
        return instance

class A(LightningDataModule, metaclass=_InitMeta):
    def __init__(self, *args):
        super().__init__()  # Call parent without kwargs
        self.a = 1
        self.b = 2

    def print_ab(self):
        print("in print ab")
        print(f"a={self.a}, b={self.b}")
        if hasattr(self, "_hparams_dict"):
            print("Hyperparameters:", self._hparams_dict)

    def _after_init(self):
        """Called only once after full initialization."""
        if hasattr(self, "_hparams_dict"):
            # Use the internal method to set hyperparameters
            self._set_hparams(self._hparams_dict)
        self.print_ab()

class B(A):
    def __init__(self, *args):
        super().__init__(*args)
        self.a += 1
        self.b += 2

class C(B):
    def __init__(self, *args):
        super().__init__(*args)
        self.a += 1
        self.b += 2

if __name__ == "__main__":
    print("Creating C instance:")
    c = C(learning_rate=0.01, batch_size=32)  # Should print 3, 6 only once
    print("\nCreating B instance:")
    b = B(learning_rate=0.02, batch_size=64)  # Should print 2, 4 only once
    print("\nCreating A instance:")
    a = A(learning_rate=0.03, batch_size=128)  # Should print 1, 2 only once```