Repository: Lightning-AI/pytorch-lightning
Pull request: LightningCLI save config multifile option required to support save to fsspec filesystem
## What does this PR do?

jsonargparse now supports saving config to fsspec filesystem. This is a necessary change for it to work by default.

Fixes #8843

### Does your PR introduce any breaking changes? If yes, please list them.

<!-- FILL IN or None -->

## Before submitting

- [x] Was this **discussed/approved** via a GitHub issue? (not for typos and docs)
- [x] Did you read the [contributor guideline](https://github.com/PyTorchLightning/pytorch-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?
- [x] 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?
- [x] Did you list all the **breaking changes** introduced by this pull request?
- [x] Did you **update the [CHANGELOG](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/CHANGELOG.md)**? (not for typos, docs, test updates, or internal minor changes/refactorings)

<!-- In the CHANGELOG, separate each item in the unreleased section by a blank line to reduce collisions -->

## PR review

Anyone in the community is welcome to review the PR.
Before you start reviewing make sure you have read [Review guidelines](https://github.com/PyTorchLightning/pytorch-lightning/wiki/Review-guidelines). In short, see the following bullet-list:

- [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

## Did you have fun?

Make sure you had fun coding 🙃


Problem statement:
SaveConfigCallback errors out when default_root_dir is an S3 path
## 🐛 Bug

Hi, when running PL Lightning using an S3 `default_root_dir` in the Trainer (while using the `LightningCLI`), I’m getting the following error:
```
Traceback (most recent call last):
File "train.py", line 56, in <module>
trainer_defaults=trainer_defaults
File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/utilities/cli.py", line 289, in __init__
self.fit()
File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/utilities/cli.py", line 432, in fit
self.trainer.fit(**self.fit_kwargs)
File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py", line 553, in fit
self._run(model)
File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py", line 865, in _run
self._call_setup_hook(model) # allow user to setup lightning_module in accelerator environment
File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py", line 1170, in _call_setup_hook
self.setup(model, stage=fn)
File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/trainer/callback_hook.py", line 52, in setup
callback.setup(self, model, stage=stage)
File "/opt/conda/lib/python3.6/site-packages/pytorch_lightning/utilities/cli.py", line 191, in setup
self.parser.save(self.config, config_path, skip_none=False, overwrite=self.overwrite)
File "/opt/conda/lib/python3.6/site-packages/jsonargparse/core.py", line 811, in save
path_fc = Path(path, mode='fc')
File "/opt/conda/lib/python3.6/site-packages/jsonargparse/util.py", line 446, in __init__
raise TypeError(ptype+' is not creatable since parent directory does not exist: '+abs_path)
TypeError: File is not creatable since parent directory does not exist: /app/s3://some_s3_path/lightning_logs/version_0/config.yaml
```
This error goes away if I add `save_config_callback=None` to the `LightningCLI` , so I think this may be a bug with the CLI saving the config callback, when the Trainer’s `default_root_dir`  points to an S3 path.


### To Reproduce
Any training script that uses `LightningCLI` and an S3 based `default_root_dir` will cause the `SaveConfigCallback` to fail.

### Expected behavior

The `SaveConfigCallback` object works with S3 based Trainer `default_root_dir` values.

Hints from issue discussion:
strange, in theory this should work since we have a makedirs call here with fsspec: 
https://github.com/PyTorchLightning/pytorch-lightning/blob/b47e3ab7ce171ced7b9d807bbca6062cd6b5c97e/pytorch_lightning/utilities/cli.py#L187

I believe you need to install this extension: 

```
pip install fsspec[s3]
```


Can you check what gets returned when you do

```python
from pytorch_lightning.utilities.cloud_io import get_filesystem
fs = get_filesystem("path/to/your/s3/folder")
```
?
Hi @awaelchli , thanks for your help and sorry for the delay in my response! The S3 integration in general works (e.g., model check-pointing to S3, S3 root dirs also work without the SaveConfig callback, etc.).

In the below code excerpt from Lightning, `config_path` and `log_path` are both valid and correct/expected S3 paths at runtime (which I confirmed with `print` statements):

```python
if trainer.is_global_zero:
    # save only on rank zero to avoid race conditions on DDP.
    # the `log_dir` needs to be created as we rely on the logger to do it usually
    # but it hasn't logged anything at this point
    get_filesystem(log_dir).makedirs(log_dir, exist_ok=True)
    self.parser.save(self.config, config_path, skip_none=False, overwrite=self.overwrite)
```
However, I think the issue arises because in the last line (`self.parser.save(...)`), `self.parser` is a `LightningArgumentParser` object that extends from `jsonargparse`'s `ArgumentParser`, but doesn't override its `save` method (see https://github.com/omni-us/jsonargparse/blob/45996d88e9026e345fd10d4f6bf46859cc243ca3/jsonargparse/core.py#L781). And since the `jsonargparse`'s `ArgumentParser`'s save method assumes a local filesystem, it is unable to write to S3.
@thecooltechguy I don't think we should override jsonargparse's save method. 

@carmocca @mauvilsa what do you say?

a) override their save function
b) open a feature request in jsonargparse
c) not support it

or something else.
Since `jsonargparse` already supports reading config files off s3, I think it's fair for it to implement writing to s3 too.

https://github.com/omni-us/jsonargparse/blob/df466645f64a4fdd5983197ca6fe28ca38126d04/jsonargparse/optionals.py#L112-L117

Otherwise, it could refactor the `parser.save` function so we don't have to override it entirely, just the data writing parts. Then it would be easy for us to override those as we already subclass their `ArgumentParser`

https://github.com/PyTorchLightning/pytorch-lightning/blob/13e64e6a80bc19308e99ff6aedb95ad68968b5b8/pytorch_lightning/utilities/cli.py#L37
I think best if it is added to jsonargparse.
Awesome, sounds good! Thanks!
@thecooltechguy could you please try out #9073 with the current master version of jsonargparse.

Commit message: Merge branch 'master' into cli_multifile_option

Patch to review:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 169b63ed30173..2c223ef674e3a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -187,6 +187,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - `Trainer.request_dataloader` now takes a `RunningStage` enum instance ([#8858](https://github.com/PyTorchLightning/pytorch-lightning/pull/8858))
 
 
+- `LightningCLI` by default support saving to fsspec filesystem ([#9073](https://github.com/PyTorchLightning/pytorch-lightning/pull/9073))
+
+
 - Changed `rank_zero_warn` to `NotImplementedError` in the `{train, val, test, predict}_dataloader` hooks that `Lightning(Data)Module` uses ([#9161](https://github.com/PyTorchLightning/pytorch-lightning/pull/9161))
 
 
diff --git a/pytorch_lightning/utilities/cli.py b/pytorch_lightning/utilities/cli.py
index 1d437e69a9ef0..62a86731c50ed 100644
--- a/pytorch_lightning/utilities/cli.py
+++ b/pytorch_lightning/utilities/cli.py
@@ -160,11 +160,13 @@ def __init__(
         config: Union[Namespace, Dict[str, Any]],
         config_filename: str,
         overwrite: bool = False,
+        multifile: bool = False,
     ) -> None:
         self.parser = parser
         self.config = config
         self.config_filename = config_filename
         self.overwrite = overwrite
+        self.multifile = multifile
 
     def setup(self, trainer: Trainer, pl_module: LightningModule, stage: Optional[str] = None) -> None:
         # save the config in `setup` because (1) we want it to save regardless of the trainer function run
@@ -184,7 +186,9 @@ def setup(self, trainer: Trainer, pl_module: LightningModule, stage: Optional[st
             # the `log_dir` needs to be created as we rely on the logger to do it usually
             # but it hasn't logged anything at this point
             get_filesystem(log_dir).makedirs(log_dir, exist_ok=True)
-            self.parser.save(self.config, config_path, skip_none=False, overwrite=self.overwrite)
+            self.parser.save(
+                self.config, config_path, skip_none=False, overwrite=self.overwrite, multifile=self.multifile
+            )
 
     def __reduce__(self) -> Tuple[Type["SaveConfigCallback"], Tuple, Dict]:
         # `ArgumentParser` is un-pickleable. Drop it
@@ -201,6 +205,7 @@ def __init__(
         save_config_callback: Optional[Type[SaveConfigCallback]] = SaveConfigCallback,
         save_config_filename: str = "config.yaml",
         save_config_overwrite: bool = False,
+        save_config_multifile: bool = False,
         trainer_class: Union[Type[Trainer], Callable[..., Trainer]] = Trainer,
         trainer_defaults: Optional[Dict[str, Any]] = None,
         seed_everything_default: Optional[int] = None,
@@ -232,6 +237,7 @@ def __init__(
             save_config_callback: A callback class to save the training config.
             save_config_filename: Filename for the config file.
             save_config_overwrite: Whether to overwrite an existing config file.
+            save_config_multifile: When input is multiple config files, saved config preserves this structure.
             trainer_class: An optional subclass of the :class:`~pytorch_lightning.trainer.trainer.Trainer` class or a
                 callable which returns a :class:`~pytorch_lightning.trainer.trainer.Trainer` instance when called.
             trainer_defaults: Set to override Trainer defaults or add persistent callbacks.
@@ -255,6 +261,7 @@ def __init__(
         self.save_config_callback = save_config_callback
         self.save_config_filename = save_config_filename
         self.save_config_overwrite = save_config_overwrite
+        self.save_config_multifile = save_config_multifile
         self.trainer_class = trainer_class
         self.trainer_defaults = trainer_defaults or {}
         self.seed_everything_default = seed_everything_default
@@ -421,7 +428,11 @@ def _instantiate_trainer(self, config: Dict[str, Any], callbacks: List[Callback]
                 config["callbacks"].append(self.trainer_defaults["callbacks"])
         if self.save_config_callback and not config["fast_dev_run"]:
             config_callback = self.save_config_callback(
-                self.parser, self.config, self.save_config_filename, overwrite=self.save_config_overwrite
+                self.parser,
+                self.config,
+                self.save_config_filename,
+                overwrite=self.save_config_overwrite,
+                multifile=self.save_config_multifile,
             )
             config["callbacks"].append(config_callback)
         return self.trainer_class(**config)


Code review comments:
diff_hunk: @@ -187,6 +187,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - `Trainer.request_dataloader` now takes a `RunningStage` enum instance ([#8858](https://github.com/PyTorchLightning/pytorch-lightning/pull/8858))
 
 
+- `LightningCLI` by default support saving to fsspec filesystem ([#9073](https://github.com/PyTorchLightning/pytorch-lightning/pull/9073))
path: CHANGELOG.md
text: @user1:
what does default mean here?

@author:
It means that if the default were `multifile: bool = True` and someone has the `default_root_dir` an fsspec then it will give a message that this is not supported. The multifile option is for when as input config files within other config files are given, e.g. a general `config.yaml` which internally references another file say `model.yaml`. If multifile is true, saving the config generates the separate config files.

@user2:
I think this should be in the "Added" section (not "Changed") and be:

```suggestion
- Added `LightningCLI` support for saving multifile configuration files ([#9073](https://github.com/PyTorchLightning/pytorch-lightning/pull/9073))
```

@author:
Moved it to Added section and changed the text to make it more clear