test_optimizers.py

# Copyright The Lightning AI team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# 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 unittest import mock
from unittest.mock import call

import pytest
import torch
from lightning.pytorch import Trainer
from lightning.pytorch.callbacks import ModelCheckpoint
from lightning.pytorch.core.optimizer import (
    _configure_optimizers,
    _configure_schedulers_automatic_opt,
    _init_optimizers_and_lr_schedulers,
)
from lightning.pytorch.demos.boring_classes import BoringDataModule, BoringModel
from lightning.pytorch.utilities.exceptions import MisconfigurationException
from lightning.pytorch.utilities.types import LRSchedulerConfig
from torch import optim

from tests_pytorch.helpers.runif import RunIf


def test_optimizer_with_scheduling(tmp_path):
    """Verify that learning rate scheduling is working."""
    model = BoringModel()
    trainer = Trainer(
        default_root_dir=tmp_path, max_epochs=1, limit_val_batches=0.1, limit_train_batches=0.2, val_check_interval=0.5
    )
    trainer.fit(model)

    init_lr = 0.1
    adjusted_lr = [pg["lr"] for pg in trainer.optimizers[0].param_groups]

    assert len(trainer.lr_scheduler_configs) == 1
    assert all(a == adjusted_lr[0] for a in adjusted_lr)
    assert init_lr * 0.1 == adjusted_lr[0]


def test_multi_optimizer_with_scheduling(tmp_path):
    """Verify that learning rate scheduling is working."""

    class Model(BoringModel):
        init_lr = 5e-4

        def training_step(self, batch, batch_idx):
            opt1, opt2 = self.optimizers()
            loss = self.loss(self.step(batch))
            opt1.zero_grad()
            opt2.zero_grad()
            self.manual_backward(loss)
            opt1.step()
            opt2.step()

        def on_train_epoch_end(self):
            scheduler1, scheduler2 = self.lr_schedulers()
            scheduler1.step()
            scheduler2.step()

        def configure_optimizers(self):
            optimizer1 = optim.Adam(self.parameters(), lr=self.init_lr)
            optimizer2 = optim.Adam(self.parameters(), lr=self.init_lr)
            lr_scheduler1 = optim.lr_scheduler.StepLR(optimizer1, step_size=1)
            lr_scheduler2 = optim.lr_scheduler.StepLR(optimizer2, step_size=1)
            return [optimizer1, optimizer2], [lr_scheduler1, lr_scheduler2]

    model = Model()
    model.automatic_optimization = False
    trainer = Trainer(default_root_dir=tmp_path, max_epochs=1, limit_val_batches=0.1, limit_train_batches=0.2)
    trainer.fit(model)

    adjusted_lr1 = [pg["lr"] for pg in trainer.optimizers[0].param_groups]
    adjusted_lr2 = [pg["lr"] for pg in trainer.optimizers[1].param_groups]

    assert len(trainer.lr_scheduler_configs) == 2
    assert all(a == adjusted_lr1[0] for a in adjusted_lr1)
    assert all(a == adjusted_lr2[0] for a in adjusted_lr2)
    assert model.init_lr * 0.1 == adjusted_lr1[0]
    assert model.init_lr * 0.1 == adjusted_lr2[0]


def test_reducelronplateau_with_no_monitor_raises(tmp_path):
    """Test exception when a ReduceLROnPlateau is used with no monitor."""
    model = BoringModel()
    optimizer = optim.Adam(model.parameters())
    model.configure_optimizers = lambda: ([optimizer], [optim.lr_scheduler.ReduceLROnPlateau(optimizer)])
    trainer = Trainer(default_root_dir=tmp_path, fast_dev_run=True)
    with pytest.raises(
        MisconfigurationException, match="`configure_optimizers` must include a monitor when a `ReduceLROnPlateau`"
    ):
        trainer.fit(model)


def test_reducelronplateau_with_no_monitor_in_lr_scheduler_dict_raises(tmp_path):
    """Test exception when lr_scheduler dict has a ReduceLROnPlateau with no monitor."""
    model = BoringModel()
    optimizer = optim.Adam(model.parameters())
    model.configure_optimizers = lambda: {
        "optimizer": optimizer,
        "lr_scheduler": {"scheduler": optim.lr_scheduler.ReduceLROnPlateau(optimizer)},
    }
    trainer = Trainer(default_root_dir=tmp_path, fast_dev_run=True)
    with pytest.raises(MisconfigurationException, match="must include a monitor when a `ReduceLROnPlateau`"):
        trainer.fit(model)


def test_onecyclelr_with_epoch_interval_warns():
    """Test warning when a OneCycleLR is used and interval is epoch."""
    model = BoringModel()
    optimizer = optim.Adam(model.parameters())
    lr_scheduler = {"scheduler": optim.lr_scheduler.OneCycleLR(optimizer, max_lr=0.01, total_steps=3)}
    with pytest.warns(RuntimeWarning, match="Are you sure you didn't mean 'interval': 'step'?"):
        _configure_schedulers_automatic_opt([lr_scheduler], None)


def test_scheduler_initialized_with_custom_reduceonplateau():
    """Test for initialize custom scheduler with `reduce_on_plateau` argument."""

    class CustomReduceLROnPlateau:
        pass

    lr_scheduler = {"reduce_on_plateau": True, "scheduler": CustomReduceLROnPlateau(), "monitor": "my_loss"}
    config = _configure_schedulers_automatic_opt([lr_scheduler], None)
    assert isinstance(config[0].scheduler, CustomReduceLROnPlateau)
    assert config[0].reduce_on_plateau


def test_reducelronplateau_scheduling(tmp_path):
    class TestModel(BoringModel):
        def training_step(self, batch, batch_idx):
            self.log("foo", batch_idx)
            return super().training_step(batch, batch_idx)

        def configure_optimizers(self):
            optimizer = optim.Adam(self.parameters())
            return {
                "optimizer": optimizer,
                "lr_scheduler": optim.lr_scheduler.ReduceLROnPlateau(optimizer),
                "monitor": "foo",
            }

    model = TestModel()
    trainer = Trainer(default_root_dir=tmp_path, fast_dev_run=True)
    trainer.fit(model)

    lr_scheduler = trainer.lr_scheduler_configs[0]
    assert lr_scheduler == LRSchedulerConfig(
        scheduler=lr_scheduler.scheduler,
        monitor="foo",
        interval="epoch",
        frequency=1,
        reduce_on_plateau=True,
        strict=True,
        name=None,
    )


def test_optimizer_return_options(tmp_path):
    trainer = Trainer(default_root_dir=tmp_path)
    model = BoringModel()
    trainer.strategy.connect(model)
    trainer.lightning_module.trainer = trainer

    # single optimizer
    opt_a = optim.Adam(model.parameters(), lr=0.002)
    opt_b = optim.SGD(model.parameters(), lr=0.002)
    scheduler_a = optim.lr_scheduler.StepLR(opt_a, 10)
    optim.lr_scheduler.StepLR(opt_b, 10)

    # single optimizer
    model.configure_optimizers = lambda: opt_a
    opt, lr_sched = _init_optimizers_and_lr_schedulers(model)
    assert len(opt) == 1
    assert len(lr_sched) == 0

    # opt tuple
    model.automatic_optimization = False
    model.configure_optimizers = lambda: (opt_a, opt_b)
    opt, lr_sched = _init_optimizers_and_lr_schedulers(model)
    assert opt == [opt_a, opt_b]
    assert len(lr_sched) == 0

    # opt list
    model.automatic_optimization = False
    model.configure_optimizers = lambda: [opt_a, opt_b]
    opt, lr_sched = _init_optimizers_and_lr_schedulers(model)
    assert opt == [opt_a, opt_b]
    assert len(lr_sched) == 0

    ref_lr_sched = LRSchedulerConfig(
        scheduler=scheduler_a,
        interval="epoch",
        frequency=1,
        reduce_on_plateau=False,
        monitor=None,
        strict=True,
        name=None,
    )

    # opt tuple of 2 lists
    model.automatic_optimization = True
    model.configure_optimizers = lambda: ([opt_a], [scheduler_a])
    opt, lr_sched = _init_optimizers_and_lr_schedulers(model)
    assert len(opt) == len(lr_sched) == 1
    assert opt[0] == opt_a
    assert lr_sched[0] == ref_lr_sched

    # opt tuple of 1 list
    model.automatic_optimization = True
    model.configure_optimizers = lambda: ([opt_a], scheduler_a)
    opt, lr_sched = _init_optimizers_and_lr_schedulers(model)
    assert len(opt) == len(lr_sched) == 1
    assert opt[0] == opt_a
    assert lr_sched[0] == ref_lr_sched

    # opt single dictionary
    model.automatic_optimization = True
    model.configure_optimizers = lambda: {"optimizer": opt_a, "lr_scheduler": scheduler_a}
    opt, lr_sched = _init_optimizers_and_lr_schedulers(model)
    assert len(opt) == len(lr_sched) == 1
    assert opt[0] == opt_a
    assert lr_sched[0] == ref_lr_sched

    # opt list of dictionaries
    model.automatic_optimization = False
    model.configure_optimizers = lambda: [
        {"optimizer": opt_a, "lr_scheduler": scheduler_a},
        {"optimizer": opt_b, "lr_scheduler": scheduler_a},
    ]
    opt, lr_sched = _init_optimizers_and_lr_schedulers(model)
    assert len(opt) == len(lr_sched) == 2
    assert opt == [opt_a, opt_b]
    assert lr_sched == [ref_lr_sched, ref_lr_sched]


def test_none_optimizer(tmp_path):
    model = BoringModel()
    model.configure_optimizers = lambda: None
    trainer = Trainer(default_root_dir=tmp_path, max_epochs=1, limit_val_batches=0.1, limit_train_batches=0.2)
    with pytest.warns(UserWarning, match="will run with no optimizer"):
        trainer.fit(model)


def test_configure_optimizer_from_dict(tmp_path):
    """Tests if `configure_optimizer` method could return a dictionary with `optimizer` field only."""

    class TestModel(BoringModel):
        def configure_optimizers(self):
            return {"optimizer": optim.SGD(params=self.parameters(), lr=1e-03)}

    model = TestModel()
    trainer = Trainer(default_root_dir=tmp_path, fast_dev_run=True)
    trainer.fit(model)


@pytest.mark.parametrize("fn", ["validate", "test", "predict"])
def test_init_optimizers_during_evaluation_and_prediction(tmp_path, fn):
    """Test that optimizers is an empty list during evaluation and prediction."""

    class TestModel(BoringModel):
        def configure_optimizers(self):
            optimizer1 = optim.Adam(self.parameters(), lr=0.1)
            optimizer2 = optim.Adam(self.parameters(), lr=0.1)
            lr_scheduler1 = optim.lr_scheduler.StepLR(optimizer1, step_size=1)
            lr_scheduler2 = optim.lr_scheduler.StepLR(optimizer2, step_size=1)
            return [optimizer1, optimizer2], [lr_scheduler1, lr_scheduler2]

    trainer = Trainer(default_root_dir=tmp_path, fast_dev_run=2)
    train_fn = getattr(trainer, fn)
    train_fn(TestModel(), datamodule=BoringDataModule(), ckpt_path=None)

    assert len(trainer.lr_scheduler_configs) == 0
    assert len(trainer.optimizers) == 0


@pytest.mark.parametrize("complete_epoch", [True, False])
@mock.patch("torch.optim.lr_scheduler.ReduceLROnPlateau.step")
def test_lr_scheduler_strict(step_mock, tmp_path, complete_epoch):
    """Test "strict" support in lr_scheduler dict."""
    model = BoringModel()
    optimizer = optim.Adam(model.parameters())
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer)
    max_epochs = 1 if complete_epoch else None
    max_steps = -1 if complete_epoch else 1
    trainer = Trainer(default_root_dir=tmp_path, max_epochs=max_epochs, max_steps=max_steps)

    model.configure_optimizers = lambda: {
        "optimizer": optimizer,
        "lr_scheduler": {"scheduler": scheduler, "monitor": "giraffe", "strict": True},
    }

    if complete_epoch:
        with pytest.raises(
            MisconfigurationException,
            match=r"ReduceLROnPlateau conditioned on metric .* which is not available\. Available metrics are:",
        ):
            trainer.fit(model)
    else:
        trainer.fit(model)

    step_mock.assert_not_called()

    model.configure_optimizers = lambda: {
        "optimizer": optimizer,
        "lr_scheduler": {"scheduler": scheduler, "monitor": "giraffe", "strict": False},
    }

    if complete_epoch:
        trainer = Trainer(default_root_dir=tmp_path, max_epochs=max_epochs, max_steps=max_steps)
        with pytest.warns(
            RuntimeWarning, match=r"ReduceLROnPlateau conditioned on metric .* which is not available but strict"
        ):
            trainer.fit(model)

    step_mock.assert_not_called()


def test_unknown_configure_optimizers_raises(tmp_path):
    """Test exception with an unsupported configure_optimizers return."""
    model = BoringModel()
    model.configure_optimizers = lambda: 1
    trainer = Trainer(default_root_dir=tmp_path, fast_dev_run=True)
    with pytest.raises(MisconfigurationException, match="Unknown configuration for model optimizers"):
        trainer.fit(model)


def test_optimizer_config_dict_with_extra_keys_warns(tmp_path):
    """Test exception when optimizer configuration dict has extra keys."""
    model = BoringModel()
    optimizer = optim.Adam(model.parameters())
    optim_conf = {
        "optimizer": optimizer,
        "lr_scheduler": {"scheduler": optim.lr_scheduler.StepLR(optimizer, 1)},
        "foo": 1,
        "bar": 2,
    }
    with pytest.warns(RuntimeWarning, match=r"Found unsupported keys in the optimizer configuration: \{.+\}"):
        _configure_optimizers(optim_conf)


def test_multiple_optimizer_config_dicts_with_extra_keys_warns(tmp_path):
    """Test exception when multiple optimizer configuration dicts have extra keys."""
    model = BoringModel()
    optimizer1 = optim.Adam(model.parameters(), lr=0.01)
    optimizer2 = optim.Adam(model.parameters(), lr=0.01)
    lr_scheduler_config_1 = {"scheduler": optim.lr_scheduler.StepLR(optimizer1, 1)}
    lr_scheduler_config_2 = {"scheduler": optim.lr_scheduler.StepLR(optimizer2, 1)}
    optim_conf = [
        {"optimizer": optimizer1, "lr_scheduler": lr_scheduler_config_1, "foo": 1, "bar": 2},
        {"optimizer": optimizer2, "lr_scheduler": lr_scheduler_config_2, "foo": 1, "bar": 2},
    ]
    with pytest.warns(RuntimeWarning, match=r"Found unsupported keys in the optimizer configuration: \{.+\}"):
        _configure_optimizers(optim_conf)


def test_lr_scheduler_with_unknown_interval_raises(tmp_path):
    """Test exception when lr_scheduler dict has unknown interval param value."""
    model = BoringModel()
    optimizer = optim.Adam(model.parameters())
    model.configure_optimizers = lambda: {
        "optimizer": optimizer,
        "lr_scheduler": {"scheduler": optim.lr_scheduler.StepLR(optimizer, 1), "interval": "incorrect_unknown_value"},
    }
    trainer = Trainer(default_root_dir=tmp_path, fast_dev_run=True)
    with pytest.raises(MisconfigurationException, match=r'The "interval" key in lr scheduler dict must be'):
        trainer.fit(model)


def test_lr_scheduler_with_extra_keys_warns(tmp_path):
    """Test warning when lr_scheduler dict has extra keys."""
    model = BoringModel()
    optimizer = optim.Adam(model.parameters())
    model.configure_optimizers = lambda: {
        "optimizer": optimizer,
        "lr_scheduler": {"scheduler": optim.lr_scheduler.StepLR(optimizer, 1), "foo": 1, "bar": 2},
    }
    trainer = Trainer(default_root_dir=tmp_path, fast_dev_run=True)
    with pytest.warns(RuntimeWarning, match=r"Found unsupported keys in the lr scheduler dict: \{.+\}"):
        trainer.fit(model)


def test_lr_scheduler_with_no_actual_scheduler_raises(tmp_path):
    """Test exception when lr_scheduler dict has no scheduler."""
    model = BoringModel()
    model.configure_optimizers = lambda: {"optimizer": optim.Adam(model.parameters()), "lr_scheduler": {}}
    trainer = Trainer(default_root_dir=tmp_path, fast_dev_run=True)
    with pytest.raises(MisconfigurationException, match='The lr scheduler dict must have the key "scheduler"'):
        trainer.fit(model)


def test_invalid_optimizer_in_scheduler(tmp_path):
    """Test exception when optimizer attached to lr_schedulers wasn't returned."""

    class InvalidOptimizerModel(BoringModel):
        def configure_optimizers(self):
            opt1 = optim.SGD(self.layer.parameters(), lr=0.1)
            opt2 = optim.SGD(self.layer.parameters(), lr=0.1)
            lr_scheduler = optim.lr_scheduler.StepLR(opt2, step_size=1)
            return [opt1], [lr_scheduler]

    model = InvalidOptimizerModel()
    trainer = Trainer(default_root_dir=tmp_path, fast_dev_run=True)
    with pytest.raises(MisconfigurationException, match="attached with an optimizer that wasn't returned"):
        trainer.fit(model)


def test_invalid_optimizer_dict_raises(tmp_path):
    """Test exception when lr_scheduler dict has no scheduler."""

    class DummyModel(BoringModel):
        def configure_optimizers(self):
            return [{"optimizer": optim.Adam(self.parameters())}, optim.Adam(self.parameters())]

    model = DummyModel()
    trainer = Trainer(default_root_dir=tmp_path, fast_dev_run=True)
    with pytest.raises(MisconfigurationException, match="Unknown configuration for model optimizers"):
        trainer.fit(model)


@RunIf(min_cuda_gpus=2, standalone=True)
def test_optimizer_state_on_device(tmp_path):
    """Test that optimizers that create state initially at instantiation still end up with the state on the GPU."""

    class TestModel(BoringModel):
        def configure_optimizers(self):
            # Adagrad creates state tensors immediately, model is not yet on GPU.
            return optim.Adagrad(self.parameters())

        def on_train_start(self, *args, **kwargs):
            opt = self.optimizers()
            _, state = next(iter(opt.state.items()))
            assert state["sum"].device == torch.device("cuda", self.local_rank) == self.device

    model = TestModel()
    trainer = Trainer(
        default_root_dir=tmp_path,
        accelerator="gpu",
        devices=2,
        strategy="ddp",
        fast_dev_run=True,
        enable_progress_bar=False,
        enable_model_summary=False,
    )
    trainer.fit(model)


@pytest.mark.parametrize("check_val_every_n_epoch", [1, 2])
@mock.patch("torch.optim.lr_scheduler.StepLR.step")
def test_lr_scheduler_epoch_step_frequency(mocked_sched, check_val_every_n_epoch, tmp_path):
    epochs = 4
    expected_steps = epochs + 1  # every LRScheduler gets called once at init

    model = BoringModel()
    trainer = Trainer(
        default_root_dir=tmp_path,
        limit_train_batches=2,
        limit_val_batches=2,
        check_val_every_n_epoch=check_val_every_n_epoch,
        max_epochs=epochs,
    )
    trainer.fit(model)
    assert mocked_sched.call_count == expected_steps


@pytest.mark.parametrize(("every_n_train_steps", "epoch_interval"), [(None, True), (2, False), (2, True)])
def test_lr_scheduler_state_updated_before_saving(tmp_path, every_n_train_steps, epoch_interval):
    batches = 2
    max_epochs = 1
    lr, gamma = 1, 10
    trainer = Trainer(
        default_root_dir=tmp_path,
        enable_progress_bar=False,
        logger=False,
        max_epochs=max_epochs,
        limit_train_batches=batches,
        limit_val_batches=1,
        callbacks=[ModelCheckpoint(dirpath=tmp_path, every_n_train_steps=every_n_train_steps)],
    )

    class TestModel(BoringModel):
        def configure_optimizers(self):
            optimizer = torch.optim.SGD(self.parameters(), lr=lr)
            lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=gamma)
            lr_scheduler_config = {"scheduler": lr_scheduler}
            if not epoch_interval:
                lr_scheduler_config["interval"] = "step"
            return [optimizer], [lr_scheduler_config]

        def on_save_checkpoint(self, checkpoint):
            lr_scheduler_config = checkpoint["lr_schedulers"][0]
            # 2 batches ran. since the lr_scheduler_config interval is `step`, the step count should be 2
            assert self.trainer.global_step == batches
            compare_to = max_epochs if epoch_interval else batches
            assert lr_scheduler_config["_step_count"] - 1 == compare_to  # step count starts at 1
            assert lr_scheduler_config["_last_lr"] == [lr * gamma**compare_to]
            self.on_save_checkpoint_called = True

    model = TestModel()
    trainer.fit(model)
    assert model.on_save_checkpoint_called


@pytest.mark.parametrize("save_on_train_epoch_end", [False, True])
def test_plateau_scheduler_lr_step_interval_updated_after_saving(tmp_path, save_on_train_epoch_end):
    batches = 4
    trainer = Trainer(
        default_root_dir=tmp_path,
        enable_progress_bar=False,
        logger=False,
        max_epochs=1,
        limit_train_batches=batches,
        limit_val_batches=1,
        callbacks=[ModelCheckpoint(dirpath=tmp_path, save_on_train_epoch_end=save_on_train_epoch_end)],
    )

    class Model(BoringModel):
        def training_step(self, batch, batch_idx):
            self.log("foo", batch_idx)
            return super().training_step(batch, batch_idx)

        def configure_optimizers(self):
            optimizer = torch.optim.Adam(self.parameters())

            lr_scheduler1 = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer)
            lr_scheduler_config_1 = {"scheduler": lr_scheduler1, "interval": "step", "monitor": "foo"}

            lr_scheduler2 = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1)
            lr_scheduler_config_2 = {"scheduler": lr_scheduler2, "interval": "step"}
            return [optimizer], [lr_scheduler_config_1, lr_scheduler_config_2]

        def on_save_checkpoint(self, checkpoint):
            lr_scheduler_config_1 = checkpoint["lr_schedulers"][0]
            last_epoch = lr_scheduler_config_1["last_epoch"]
            assert last_epoch == batches - (not save_on_train_epoch_end)  # last epoch starts at 0

            lr_scheduler_config_2 = checkpoint["lr_schedulers"][1]
            assert lr_scheduler_config_2["_step_count"] - 1 == batches  # step count starts at 1

            self.on_save_checkpoint_called = True

    model = Model()
    trainer.fit(model)
    assert model.on_save_checkpoint_called


def test_lr_scheduler_step_hook(tmp_path):
    """Test that custom lr scheduler works and `lr_scheduler_step` is called at appropriate time."""

    class CustomEpochScheduler:
        def __init__(self, optimizer):
            self.optimizer = optimizer

        def step(self, epoch): ...

        def state_dict(self): ...

        def load_state_dict(self, state_dict): ...

    class CustomBoringModel(BoringModel):
        def lr_scheduler_step(self, scheduler: int, metric):
            # step-level
            if isinstance(scheduler, torch.optim.lr_scheduler.StepLR):
                super().lr_scheduler_step(scheduler, metric)
            # epoch-level, custom scheduler
            elif isinstance(scheduler, CustomEpochScheduler):
                scheduler.step(epoch=self.current_epoch)

        def configure_optimizers(self):
            optimizer = torch.optim.SGD(self.layer.parameters(), lr=1e-2)
            lr_scheduler1 = {"scheduler": torch.optim.lr_scheduler.StepLR(optimizer, step_size=1), "interval": "step"}
            lr_scheduler2 = CustomEpochScheduler(optimizer)
            return [optimizer], [lr_scheduler1, lr_scheduler2]

    model = CustomBoringModel()
    max_epochs = 3
    limit_train_batches = 2
    trainer = Trainer(
        default_root_dir=tmp_path,
        enable_checkpointing=False,
        logger=False,
        max_epochs=max_epochs,
        limit_train_batches=limit_train_batches,
        limit_val_batches=0,
    )
    with mock.patch.object(CustomEpochScheduler, "step") as mock_method_epoch, mock.patch.object(
        torch.optim.lr_scheduler.StepLR, "step"
    ) as mock_method_step:
        trainer.fit(model)

    assert mock_method_epoch.mock_calls == [call(epoch=e) for e in range(max_epochs)]
    # first step is called by PyTorch LRScheduler
    assert mock_method_step.call_count == max_epochs * limit_train_batches + 1


def test_invalid_scheduler_missing_state_dict():
    """Test that custom lr scheduler raises an error if it's missing the state dict."""

    class CustomScheduler:
        def __init__(self, optimizer):
            self.optimizer = optimizer

        def step(self): ...

    class CustomBoringModel(BoringModel):
        def configure_optimizers(self):
            opt = torch.optim.SGD(self.parameters(), lr=1e-2)
            lr_scheduler = CustomScheduler(opt)
            return {"optimizer": opt, "lr_scheduler": lr_scheduler}

    model = CustomBoringModel()
    model.trainer = Trainer()
    with pytest.raises(TypeError, match="provided lr scheduler `CustomScheduler` is invalid"):
        _init_optimizers_and_lr_schedulers(model)


@pytest.mark.parametrize("override", [False, True])
def test_invalid_lr_scheduler_with_custom_step_method(override):
    """Test that custom lr scheduler raises an error if it doesn't follow PyTorch LR Scheduler API."""

    class CustomScheduler:
        def __init__(self, optimizer):
            self.optimizer = optimizer

        def step(self, foobar):  # breaks the API, forces user to override `lr_scheduler_step`
            ...

        def state_dict(self): ...

        def load_state_dict(self, state_dict): ...

    class CustomBoringModel(BoringModel):
        def configure_optimizers(self):
            opt = torch.optim.SGD(self.parameters(), lr=1e-2)
            lr_scheduler = CustomScheduler(opt)
            return {"optimizer": opt, "lr_scheduler": lr_scheduler}

    model = CustomBoringModel()
    model.trainer = Trainer()
    if override:

        def lr_scheduler_step(*_): ...

        # the user did override the hook, no error
        model.lr_scheduler_step = lr_scheduler_step
        _init_optimizers_and_lr_schedulers(model)
    else:
        with pytest.raises(MisconfigurationException, match="CustomScheduler` doesn't follow"):
            _init_optimizers_and_lr_schedulers(model)