[CLI] Save only the configuration used (#11532)

Author: carmocca

CHANGELOG.md

@@ -11,12 +11,11 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Pin sphinx-autodoc-typehints with <v1.15 ([#11400](https://github.com/PyTorchLightning/pytorch-lightning/pull/11400))
 - Skip testing with PyTorch 1.7 and Python 3.9 on Ubuntu ([#11217](https://github.com/PyTorchLightning/pytorch-lightning/pull/11217))
 - Fixed type promotion when tensors of higher category than float are logged ([#11401](https://github.com/PyTorchLightning/pytorch-lightning/pull/11401))
+- Fixed the format of the configuration saved automatically by the CLI's `SaveConfigCallback` ([#11532](https://github.com/PyTorchLightning/pytorch-lightning/pull/11532))
 
 ### Changed
 
 - Changed `LSFEnvironment` to use `LSB_DJOB_RANKFILE` environment variable instead of `LSB_HOSTS` for determining node rank and main address ([#10825](https://github.com/PyTorchLightning/pytorch-lightning/pull/10825))
-
-
 - Disbled sampler replacement when using `IterableDataset` ([#11507](https://github.com/PyTorchLightning/pytorch-lightning/pull/11507))
 
 

pytorch_lightning/utilities/cli.py

@@ -685,8 +685,8 @@ 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._parser(self.subcommand),
+                self.config.get(str(self.subcommand), self.config),
                 self.save_config_filename,
                 overwrite=self.save_config_overwrite,
                 multifile=self.save_config_multifile,
@@ -769,9 +769,7 @@ def configure_optimizers(
 
     def _get(self, config: Dict[str, Any], key: str, default: Optional[Any] = None) -> Any:
         """Utility to get a config value which might be inside a subcommand."""
-        if self.subcommand is not None:
-            return config[self.subcommand].get(key, default)
-        return config.get(key, default)
+        return config.get(str(self.subcommand), config).get(key, default)
 
     def _run_subcommand(self, subcommand: str) -> None:
         """Run the chosen subcommand."""

tests/utilities/test_cli.py

@@ -347,9 +347,7 @@ def test_lightning_cli_args(tmpdir):
     with open(config_path) as f:
         loaded_config = yaml.safe_load(f.read())
 
-    loaded_config = loaded_config["fit"]
     cli_config = cli.config["fit"]
-
     assert cli_config["seed_everything"] == 1234
     assert "model" not in loaded_config and "model" not in cli_config  # no arguments to include
     assert loaded_config["data"] == cli_config["data"]
@@ -403,9 +401,7 @@ def test_lightning_cli_config_and_subclass_mode(tmpdir):
     with open(config_path) as f:
         loaded_config = yaml.safe_load(f.read())
 
-    loaded_config = loaded_config["fit"]
     cli_config = cli.config["fit"]
-
     assert loaded_config["model"] == cli_config["model"]
     assert loaded_config["data"] == cli_config["data"]
     assert loaded_config["trainer"] == cli_config["trainer"]
@@ -1251,6 +1247,10 @@ def test_lightning_cli_config_before_subcommand():
     test_mock.assert_called_once_with(cli.trainer, model=cli.model, verbose=True, ckpt_path="foobar")
     assert cli.trainer.limit_test_batches == 1
 
+    save_config_callback = cli.trainer.callbacks[0]
+    assert save_config_callback.config["trainer"]["limit_test_batches"] == 1
+    assert save_config_callback.parser.subcommand == "test"
+
     with mock.patch("sys.argv", ["any.py", f"--config={config}", "validate"]), mock.patch(
         "pytorch_lightning.Trainer.validate", autospec=True
     ) as validate_mock: