Repository: Lightning-Universe/lightning-flash

Issue:
 flash.Trainer.add_argparse_args(parser) doesn't register args to arg parse
## 🐛 Bug

```
    parser = ArgumentParser()
    parser = flash.Trainer.add_argparse_args(parser)
    args = parser.parse_args()
```

Does not register any of the trainer args but the following code does

```
    parser = ArgumentParser()
    parser = pl.Trainer.add_argparse_args(parser)
    args = parser.parse_args()
```

Since Flash wraps the PL trainer this should work the same way




Discussion:
This is a limitation in Lightning's current argparsing implementation.

We only inspect the signature of the class passed and its `__init__`:

https://github.com/PyTorchLightning/pytorch-lightning/blob/b1a7b7e9bf6f6b1bac34c4e2be687cff4c362663/pytorch_lightning/utilities/argparse.py#L205-L210

https://github.com/PyTorchLightning/pytorch-lightning/blob/b1a7b7e9bf6f6b1bac34c4e2be687cff4c362663/pytorch_lightning/utilities/argparse.py#L109-L134

To resolve this, we would need to traverse the object chain. This seems to be how jsonargparse does it:

https://github.com/omni-us/jsonargparse/blob/26ca953c66cb88de27cf5a50f9c6b82dead9636f/jsonargparse/signatures.py#L64.

As a workaround, the `pl.Trainer` can be used to add the argument (since `flash.Trainer` does not add any).
We can override the `add_argparse_args` function in the `flash.Trainer` to do this automatically, even though it's a hack.

Resolving patch:
diff --git a/flash/core/trainer.py b/flash/core/trainer.py
--- a/flash/core/trainer.py
+++ b/flash/core/trainer.py
@@ -12,14 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import warnings
+from argparse import ArgumentParser
 from functools import wraps
 from typing import Callable, List, Optional, Union
 
 import torch
-from pytorch_lightning import LightningDataModule, LightningModule, Trainer
+from pytorch_lightning import LightningDataModule, LightningModule
+from pytorch_lightning import Trainer as PlTrainer
 from pytorch_lightning.callbacks import BaseFinetuning
 from pytorch_lightning.utilities import rank_zero_warn
-from pytorch_lightning.utilities.argparse import get_init_arguments_and_types, parse_env_variables
+from pytorch_lightning.utilities.argparse import add_argparse_args, get_init_arguments_and_types, parse_env_variables
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from torch.utils.data import DataLoader
 
@@ -49,7 +51,7 @@ def insert_env_defaults(self, *args, **kwargs):
     return insert_env_defaults
 
 
-class Trainer(Trainer):
+class Trainer(PlTrainer):
 
     @_defaults_from_env_vars
     def __init__(self, *args, **kwargs):
@@ -172,3 +174,9 @@ def _merge_callbacks(old_callbacks: List, new_callbacks: List) -> List:
         override_types = new_callbacks_types.intersection(old_callbacks_types)
         new_callbacks.extend(c for c in old_callbacks if type(c) not in override_types)
         return new_callbacks
+
+    @classmethod
+    def add_argparse_args(cls, *args, **kwargs) -> ArgumentParser:
+        # the lightning trainer implementation does not support subclasses.
+        # context: https://github.com/PyTorchLightning/lightning-flash/issues/342#issuecomment-848892447
+        return add_argparse_args(PlTrainer, *args, **kwargs)


Test patch:
diff --git a/tests/core/test_trainer.py b/tests/core/test_trainer.py
--- a/tests/core/test_trainer.py
+++ b/tests/core/test_trainer.py
@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from argparse import ArgumentParser
 from typing import Any
 
 import pytest
@@ -110,3 +111,10 @@ def test_resolve_callbacks_override_warning(tmpdir):
     task = FinetuneClassificationTask(model, loss_fn=F.nll_loss)
     with pytest.warns(UserWarning, match="The model contains a default finetune callback"):
         trainer._resolve_callbacks(task, "test")
+
+
+def test_add_argparse_args():
+    parser = ArgumentParser()
+    parser = Trainer.add_argparse_args(parser)
+    args = parser.parse_args(['--gpus=1'])
+    assert args.gpus == 1