Fix typehint for WandbLogger's log_model argument (#18458)

Co-authored-by: Carlos Mocholí <carlossmocholi@gmail.com>
Co-authored-by: awaelchli <aedu.waelchli@gmail.com>

src/lightning/pytorch/CHANGELOG.md

@@ -225,6 +225,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed redundant `iter()` call to dataloader when checking dataloading configuration ([#18415](https://github.com/Lightning-AI/lightning/pull/18415))
 
 
+- Fixed an issue that wouldn't prevent the user to set the `log_model` parameter in `WandbLogger` via the LightningCLI ([#18458](https://github.com/Lightning-AI/lightning/pull/18458))
+
+
 ## [2.0.7] - 2023-08-14
 
 ### Added

src/lightning/pytorch/loggers/wandb.py

@@ -18,7 +18,7 @@
 import os
 from argparse import Namespace
 from pathlib import Path
-from typing import Any, Dict, List, Mapping, Optional, Union
+from typing import Any, Dict, List, Literal, Mapping, Optional, Union
 
 import torch.nn as nn
 from lightning_utilities.core.imports import RequirementCache
@@ -295,7 +295,7 @@ def __init__(
         id: Optional[str] = None,
         anonymous: Optional[bool] = None,
         project: Optional[str] = None,
-        log_model: Union[str, bool] = False,
+        log_model: Union[Literal["all"], bool] = False,
         experiment: Union[Run, RunDisabled, None] = None,
         prefix: str = "",
         checkpoint_name: Optional[str] = None,

tests/tests_pytorch/loggers/test_wandb.py

@@ -16,10 +16,12 @@
 from unittest import mock
 
 import pytest
+import yaml
 from lightning_utilities.test.warning import no_warning_call
 
 from lightning.pytorch import Trainer
 from lightning.pytorch.callbacks import ModelCheckpoint
+from lightning.pytorch.cli import LightningCLI
 from lightning.pytorch.demos.boring_classes import BoringModel
 from lightning.pytorch.loggers import WandbLogger
 from lightning.pytorch.utilities.exceptions import MisconfigurationException
@@ -497,3 +499,39 @@ def test_wandb_logger_download_artifact(wandb, tmpdir):
     WandbLogger.download_artifact("test_artifact", str(tmpdir), "model", True)
 
     wandb.Api().artifact.assert_called_once_with("test_artifact", type="model")
+
+
+@mock.patch("lightning.pytorch.loggers.wandb._WANDB_AVAILABLE", True)
+@mock.patch("lightning.pytorch.loggers.wandb.wandb", mock.Mock())
+@pytest.mark.parametrize(("log_model", "expected"), [("True", True), ("False", False), ("all", "all")])
+def test_wandb_logger_cli_integration(log_model, expected, monkeypatch, tmp_path):
+    """Test that the WandbLogger can be used with the LightningCLI."""
+    monkeypatch.chdir(tmp_path)
+
+    class InspectParsedCLI(LightningCLI):
+        def before_instantiate_classes(self):
+            assert self.config.trainer.logger.init_args.log_model == expected
+
+    # Create a config file with the log_model parameter set. This seems necessary to be able
+    # to set the init_args parameter of the logger on the CLI later on.
+    input_config = {
+        "trainer": {
+            "logger": {
+                "class_path": "lightning.pytorch.loggers.wandb.WandbLogger",
+                "init_args": {"log_model": log_model},
+            },
+        }
+    }
+    config_path = "config.yaml"
+    with open(config_path, "w") as f:
+        f.write(yaml.dump(input_config))
+
+    # Test case 1: Set the log_model parameter only via the config file.
+    with mock.patch("sys.argv", ["any.py", "--config", config_path]):
+        InspectParsedCLI(BoringModel, run=False, save_config_callback=None)
+
+    # Test case 2: Overwrite the log_model parameter via the command line.
+    wandb_cli_arg = f"--trainer.logger.init_args.log_model={log_model}"
+
+    with mock.patch("sys.argv", ["any.py", "--config", config_path, wandb_cli_arg]):
+        InspectParsedCLI(BoringModel, run=False, save_config_callback=None)