LightningCLI dataclass support broken

[leezu]

🐛 Bug

Usage of simple dataclasses with LightningCLI broke with jsonargparse 3.14+ which is required in recent version of pytorch-lightning. I'm unable to reproduce the issue with jsonargparse directly (but haven't tried extensively) thus report it here.

@dataclass
class Config:
    name: str


class Module(pl.LightningModule):
    def __init__(self, *, cfg: Config):
        super().__init__()

[...]

if __name__ == "__main__":
     LightningCLI(Module)

Triggers

  File "test.py", line 41, in <module>                                                           
    print(LightningCLI(Module))                               
  File ".local/lib/python3.8/site-packages/pytorch_lightning/utilities/cli.py", line 285, in __init__           
    self.instantiate_classes()                                                                                               
  File ".local/lib/python3.8/site-packages/pytorch_lightning/utilities/cli.py", line 338, in inst antiate_classes
    self.config_init = self.parser.instantiate_classes(self.config)                                                          
  File ".local/lib/python3.8/site-packages/jsonargparse/core.py", line 1126, in instantiate_classes         
    component.instantiate_class(component, cfg)
  File ".local/lib/python3.8/site-packages/jsonargparse/signatures.py", line 550, in group_instantiate_class
    parent[key] = group.group_class(**value)
TypeError: type object argument after ** must be a mapping, not Config

To Reproduce

from dataclasses import dataclass
import torch
from torch.utils.data import DataLoader, Dataset

import pytorch_lightning as pl
from pytorch_lightning.utilities.cli import LightningCLI


class RandomDataset(Dataset):
    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len


@dataclass
class Config:
    name: str


class Module(pl.LightningModule):
    def __init__(self, *, cfg: Config):
        super().__init__()

    def training_step(self, *args):
        pass

    def train_dataloader(self, *args):
        return DataLoader(RandomDataset(32, 64), batch_size=2)

    def configure_optimizers(self, *args):
        pass


if __name__ == "__main__":
    print(LightningCLI(Module))

Expected behavior

No crash.

Fix

diff --git a/jsonargparse/signatures.py b/jsonargparse/signatures.py
index ba936c6..d3d0c3e 100644
--- a/jsonargparse/signatures.py
+++ b/jsonargparse/signatures.py
@@ -3,6 +3,7 @@
 import inspect
 import re
 from argparse import Namespace
+import dataclasses
 from functools import wraps
 from typing import Any, Callable, List, Optional, Set, Tuple, Type, Union
 
@@ -547,6 +549,8 @@ def group_instantiate_class(group, cfg):
         parent = cfg
         key = group.dest
         assert '.' not in key
+    if dataclasses.is_dataclass(value):
+        value = dataclasses.asdict(value)
     parent[key] = group.group_class(**value)

@mauvilsa is above patch the correct fix? Or is this issue due to pytorch-lightning incorrectly calling jsonargparse?

[carmocca]

@leezu Can you report your exact jsonargparse and Lightning version? I cannot reproduce the failure locally using latest on both.

jsonargparse==3.19.0
lightning==master

Running your repro code gets me:

error: Configuration check failed :: Key "model.cfg.name" is required but not included in config object or its value is None.

Which is totally expected since it's not passing any default cfg value

[mauvilsa]

I see what the issue is. Before only subclasses were instantiated and then the instantiation of classes was added. However, since a dataclass is also a class and a group in the parser, it is trying to instantiate it. Even though the suggested fix would work, the proper fix would be that dataclasses are not considered for instantiation by instantiate_classes. I will look into it.

[leezu]

@carmocca sorry for missing the script invocation instructions:

python3 ~/test.py --model.cfg.name test --print_config > config
python3 ~/test.py --config config

[mauvilsa]

This is fixed in jsonargparse v3.19.1.