Repository: ludwig-ai/ludwig
Pull request: Fix changing parameters on plateau.
Fixes #2186 

Problem statement:
Comparing invalid types to determine reduced learning rate
Config:
```yaml
input_features:
  - name: doc
    type: text
    preprocessing:
      max_sequence_length: 4096
output_features:
  - name: cls
    type: category
    column: cls
trainer:
  epochs: 10
  early_stop: 3
  reduce_learning_rate_on_plateau: 1
  reduce_learning_rate_on_plateau_patience: 1
```

Error:
```python-traceback
  File "/home/ray/anaconda3/lib/python3.8/site-packages/ludwig/models/trainer.py", line 851, in train
    should_break = self._train_loop(
  File "/home/ray/anaconda3/lib/python3.8/site-packages/ludwig/models/trainer.py", line 1019, in _train_loop
    should_break = self.run_evaluation(
  File "/home/ray/anaconda3/lib/python3.8/site-packages/ludwig/models/trainer.py", line 652, in run_evaluation
    should_break = self.check_progress_on_validation(
  File "/home/ray/anaconda3/lib/python3.8/site-packages/ludwig/models/trainer.py", line 1177, in check_progress_on_validation
    self.reduce_learning_rate(
  File "/home/ray/anaconda3/lib/python3.8/site-packages/ludwig/models/trainer.py", line 1305, in reduce_learning_rate
    is_improved = improved(last_metric_value, progress_tracker.best_reduce_learning_rate_eval_metric)
  File "/home/ray/anaconda3/lib/python3.8/site-packages/ludwig/modules/metric_modules.py", line 476, in <lambda>
    return lambda x, y: x < y
TypeError: '<' not supported between instances of 'TrainerMetric' and 'float'
```

Hints from issue discussion:

Commit message: Fix changing parameters on plateau.

Patch to review:
diff --git a/ludwig/models/trainer.py b/ludwig/models/trainer.py
index 71954b62f09..ad1cf008eaa 100644
--- a/ludwig/models/trainer.py
+++ b/ludwig/models/trainer.py
@@ -140,15 +140,15 @@ def __init__(
         :param resume: Resume training a model that was being trained. (default: False).
         :type resume: Boolean
         :param skip_save_model: Disables saving model weights and hyperparameters each time the model improves. By
-                default Ludwig saves model weights after each epoch the validation metric (improves, but if the model is
-                really big that can be time consuming. If you do not want to keep the weights and just find out what
-                performance a model can get with a set of hyperparameters, use this parameter to skip it, but the model
-                will not be loadable later on. (default: False).
+                default Ludwig saves model weights after each round of evaluation the validation metric (improves, but
+                if the model is really big that can be time consuming. If you do not want to keep the weights and just
+                find out what performance a model can get with a set of hyperparameters, use this parameter to skip it,
+                but the model will not be loadable later on. (default: False).
         :type skip_save_model: Boolean
-        :param skip_save_progress: Disables saving progress each epoch. By default Ludwig saves weights and stats after
-                each epoch for enabling resuming of training, but if the model is really big that can be time consuming
-                and will uses twice as much space, use this parameter to skip it, but training cannot be resumed later
-                on. (default: False).
+        :param skip_save_progress: Disables saving progress each round of evaluation. By default Ludwig saves weights
+                and stats after each round of evaluation for enabling resuming of training, but if the model is really
+                big that can be time consuming and will uses twice as much space, use this parameter to skip it, but
+                training cannot be resumed later on. (default: False).
         :type skip_save_progress: Boolean
         :param skip_save_log: Disables saving TensorBoard logs. By default Ludwig saves logs for the TensorBoard, but if
                 it is not needed turning it off can slightly increase the overall speed. (default: False).
@@ -1154,11 +1154,13 @@ def check_progress_on_validation(
         # record how long its been since an improvement
         improved = get_improved_fun(validation_metric)
         validation_metrics = progress_tracker.validation_metrics[validation_output_feature_name]
-        if improved(validation_metrics[validation_metric][-1][-1], progress_tracker.best_eval_metric):
+        last_validation_metric = validation_metrics[validation_metric][-1]
+        last_validation_metric_value = last_validation_metric[-1]
+
+        if improved(last_validation_metric_value, progress_tracker.best_eval_metric):
             progress_tracker.last_improvement_steps = progress_tracker.steps
-            progress_tracker.best_eval_metric = progress_tracker.validation_metrics[validation_output_feature_name][
-                validation_metric
-            ][-1][-1]
+            progress_tracker.best_eval_metric = last_validation_metric_value
+
             if self.is_coordinator() and not skip_save_model:
                 torch.save(self.model.state_dict(), model_weights_path)
                 logger.info(
@@ -1192,10 +1194,10 @@ def check_progress_on_validation(
                 and not progress_tracker.num_reductions_learning_rate >= reduce_learning_rate_on_plateau
             ):
                 logger.info(
-                    f"Last learning rate reduction happened {progress_tracker.last_learning_rate_reduction} epoch(s) "
+                    f"Last learning rate reduction happened {progress_tracker.last_learning_rate_reduction} step(s) "
                     f"ago, improvement of {validation_output_feature_name} {reduce_learning_rate_eval_split} "
                     f"{reduce_learning_rate_eval_metric} happened "
-                    f"{progress_tracker.last_reduce_learning_rate_eval_metric_improvement} epoch(s) ago."
+                    f"{progress_tracker.last_reduce_learning_rate_eval_metric_improvement} step(s) ago."
                 )
 
         # ========== Increase Batch Size Plateau logic =========
@@ -1221,10 +1223,10 @@ def check_progress_on_validation(
             ):
                 logger.info(
                     "Last batch size increase "
-                    f"happened {progress_tracker.last_increase_batch_size} epoch(s) ago, "
+                    f"happened {progress_tracker.last_increase_batch_size} step(s) ago, "
                     f"improvement of {validation_output_feature_name} {increase_batch_size_eval_split} "
                     f"{increase_batch_size_eval_metric} happened "
-                    f"{progress_tracker.last_increase_batch_size_eval_metric_improvement} epoch(s) ago."
+                    f"{progress_tracker.last_increase_batch_size_eval_metric_improvement} step(s) ago."
                 )
 
         # ========== Early Stop logic ==========
@@ -1281,14 +1283,15 @@ def resume_weights_and_optimizer(
 
     def reduce_learning_rate(
         self,
-        progress_tracker,
-        validation_output_feature_name,
-        reduce_learning_rate_on_plateau,
-        reduce_learning_rate_on_plateau_patience,
-        reduce_learning_rate_on_plateau_rate,
-        reduce_learning_rate_eval_metric=LOSS,
-        reduce_learning_rate_eval_split=TRAINING,
+        progress_tracker: ProgressTracker,
+        validation_output_feature_name: str,
+        reduce_learning_rate_on_plateau: int,
+        reduce_learning_rate_on_plateau_patience: int,
+        reduce_learning_rate_on_plateau_rate: float,
+        reduce_learning_rate_eval_metric: str = LOSS,
+        reduce_learning_rate_eval_split: str = TRAINING,
     ):
+        """Uses the progress tracker to determine if the learning rate should be reduced."""
         if not (progress_tracker.num_reductions_learning_rate >= reduce_learning_rate_on_plateau):
 
             if reduce_learning_rate_eval_split == TRAINING:
@@ -1299,21 +1302,22 @@ def reduce_learning_rate(
                 split_metrics = progress_tracker.test_metrics
 
             validation_metric = reduce_learning_rate_eval_metric
-            last_metric_value = split_metrics[validation_output_feature_name][validation_metric][-1]
+            last_metric: TrainerMetric = split_metrics[validation_output_feature_name][validation_metric][-1]
+            last_metric_value = last_metric[-1]
 
             improved = get_improved_fun(validation_metric)
             is_improved = improved(last_metric_value, progress_tracker.best_reduce_learning_rate_eval_metric)
             if is_improved:
                 # we update the best metric value and set it to the current one
-                # and reset last improvement epoch count
+                # and reset last improvement step count
                 progress_tracker.best_reduce_learning_rate_eval_metric = last_metric_value
                 progress_tracker.last_reduce_learning_rate_eval_metric_improvement = 0
             else:
                 progress_tracker.last_reduce_learning_rate_eval_metric_improvement += 1
                 if not is_improved and (
-                    # learning rate reduction happened more than N epochs ago
+                    # learning rate reduction happened more than N steps ago
                     progress_tracker.last_learning_rate_reduction >= reduce_learning_rate_on_plateau_patience
-                    # No improvement of the evaluation metric since more than N epochs ago
+                    # No improvement of the evaluation metric since more than N steps ago
                     and progress_tracker.last_reduce_learning_rate_eval_metric_improvement
                     >= reduce_learning_rate_on_plateau_patience
                 ):
@@ -1339,15 +1343,16 @@ def reduce_learning_rate(
 
     def increase_batch_size(
         self,
-        progress_tracker,
-        validation_output_feature_name,
-        increase_batch_size_on_plateau,
-        increase_batch_size_on_plateau_patience,
-        increase_batch_size_on_plateau_rate,
-        increase_batch_size_on_plateau_max,
-        increase_batch_size_eval_metric=LOSS,
-        increase_batch_size_eval_split=TRAINING,
+        progress_tracker: ProgressTracker,
+        validation_output_feature_name: str,
+        increase_batch_size_on_plateau: int,
+        increase_batch_size_on_plateau_patience: int,
+        increase_batch_size_on_plateau_rate: float,
+        increase_batch_size_on_plateau_max: int,
+        increase_batch_size_eval_metric: str = LOSS,
+        increase_batch_size_eval_split: str = TRAINING,
     ):
+        """Uses the progress tracker to determine if the batch size should be increased."""
         if (
             not progress_tracker.num_increases_batch_size >= increase_batch_size_on_plateau
             and not progress_tracker.batch_size == increase_batch_size_on_plateau_max
@@ -1361,22 +1366,23 @@ def increase_batch_size(
                 split_metrics = progress_tracker.test_metrics
 
             validation_metric = increase_batch_size_eval_metric
-            last_metric_value = split_metrics[validation_output_feature_name][validation_metric][-1]
+            last_metric = split_metrics[validation_output_feature_name][validation_metric][-1]
+            last_metric_value = last_metric[-1]
 
             improved = get_improved_fun(validation_metric)
             is_improved = improved(last_metric_value, progress_tracker.best_increase_batch_size_eval_metric)
             if is_improved:
                 # We update the best metric value and set it to the current one, and reset last
-                # improvement epoch count
+                # improvement step count
                 progress_tracker.best_increase_batch_size_eval_metric = last_metric_value
                 progress_tracker.last_increase_batch_size_eval_metric_improvement = 0
             else:
                 progress_tracker.last_increase_batch_size_eval_metric_improvement += 1
                 if not is_improved and (
-                    # Batch size increase happened more than N epochs ago
+                    # Batch size increase happened more than N steps ago
                     progress_tracker.last_increase_batch_size >= increase_batch_size_on_plateau_patience
                     and (
-                        # No improvement of the evaluation metric since more than N epochs ago
+                        # No improvement of the evaluation metric since more than N steps ago
                         progress_tracker.last_increase_batch_size_eval_metric_improvement
                         >= increase_batch_size_on_plateau_patience
                     )
diff --git a/tests/integration_tests/test_trainer.py b/tests/integration_tests/test_trainer.py
index bd0a9b55f50..986a6adfa64 100644
--- a/tests/integration_tests/test_trainer.py
+++ b/tests/integration_tests/test_trainer.py
@@ -1,6 +1,5 @@
 import os
 import shutil
-import tempfile
 from unittest import mock
 
 import pytest
@@ -56,58 +55,57 @@ def ray_test_cluster():
 
 
 def test_tune_batch_size_and_lr(tmpdir):
-    with tempfile.TemporaryDirectory() as outdir:
-        input_features = [sequence_feature(reduce_output="sum")]
-        output_features = [category_feature(vocab_size=2, reduce_input="sum")]
+    input_features = [sequence_feature(reduce_output="sum")]
+    output_features = [category_feature(vocab_size=2, reduce_input="sum")]
 
-        csv_filename = os.path.join(tmpdir, "training.csv")
-        data_csv = generate_data(input_features, output_features, csv_filename)
-        val_csv = shutil.copyfile(data_csv, os.path.join(tmpdir, "validation.csv"))
-        test_csv = shutil.copyfile(data_csv, os.path.join(tmpdir, "test.csv"))
+    csv_filename = os.path.join(tmpdir, "training.csv")
+    data_csv = generate_data(input_features, output_features, csv_filename)
+    val_csv = shutil.copyfile(data_csv, os.path.join(tmpdir, "validation.csv"))
+    test_csv = shutil.copyfile(data_csv, os.path.join(tmpdir, "test.csv"))
 
-        config = {
-            "input_features": input_features,
-            "output_features": output_features,
-            "combiner": {"type": "concat", "output_size": 14},
-            TRAINER: {
-                "epochs": 2,
-                "batch_size": "auto",
-                "eval_batch_size": "auto",
-                "learning_rate": "auto",
-            },
-        }
+    config = {
+        "input_features": input_features,
+        "output_features": output_features,
+        "combiner": {"type": "concat", "output_size": 14},
+        TRAINER: {
+            "epochs": 2,
+            "batch_size": "auto",
+            "eval_batch_size": "auto",
+            "learning_rate": "auto",
+        },
+    }
 
-        model = LudwigModel(config, backend=LocalTestBackend())
+    model = LudwigModel(config, backend=LocalTestBackend())
 
-        # check preconditions
-        assert model.config[TRAINER][BATCH_SIZE] == "auto"
-        assert model.config[TRAINER][EVAL_BATCH_SIZE] == "auto"
-        assert model.config[TRAINER][LEARNING_RATE] == "auto"
+    # check preconditions
+    assert model.config[TRAINER][BATCH_SIZE] == "auto"
+    assert model.config[TRAINER][EVAL_BATCH_SIZE] == "auto"
+    assert model.config[TRAINER][LEARNING_RATE] == "auto"
 
-        _, _, output_directory = model.train(
-            training_set=data_csv, validation_set=val_csv, test_set=test_csv, output_directory=outdir
-        )
+    _, _, output_directory = model.train(
+        training_set=data_csv, validation_set=val_csv, test_set=test_csv, output_directory=tmpdir
+    )
 
-        def check_postconditions(model):
-            # check batch size
-            assert model.config[TRAINER][BATCH_SIZE] != "auto"
-            assert model.config[TRAINER][BATCH_SIZE] > 1
+    def check_postconditions(model):
+        # check batch size
+        assert model.config[TRAINER][BATCH_SIZE] != "auto"
+        assert model.config[TRAINER][BATCH_SIZE] > 1
 
-            assert model.config[TRAINER][EVAL_BATCH_SIZE] != "auto"
-            assert model.config[TRAINER][EVAL_BATCH_SIZE] > 1
+        assert model.config[TRAINER][EVAL_BATCH_SIZE] != "auto"
+        assert model.config[TRAINER][EVAL_BATCH_SIZE] > 1
 
-            assert model.config[TRAINER][BATCH_SIZE] == model.config[TRAINER][EVAL_BATCH_SIZE]
+        assert model.config[TRAINER][BATCH_SIZE] == model.config[TRAINER][EVAL_BATCH_SIZE]
 
-            # check learning rate
-            assert model.config[TRAINER][LEARNING_RATE] != "auto"
-            assert model.config[TRAINER][LEARNING_RATE] > 0
+        # check learning rate
+        assert model.config[TRAINER][LEARNING_RATE] != "auto"
+        assert model.config[TRAINER][LEARNING_RATE] > 0
 
-        check_postconditions(model)
+    check_postconditions(model)
 
-        model = LudwigModel.load(os.path.join(output_directory, "model"))
+    model = LudwigModel.load(os.path.join(output_directory, "model"))
 
-        # loaded model should retain the tuned params
-        check_postconditions(model)
+    # loaded model should retain the tuned params
+    check_postconditions(model)
 
 
 @pytest.mark.parametrize("learning_rate_scaling, expected_lr", [("constant", 1), ("sqrt", 2), ("linear", 4)])
@@ -137,3 +135,27 @@ def test_scale_lr(learning_rate_scaling, expected_lr, tmpdir, ray_test_cluster):
 
     actual_lr = ray.get(run_scale_lr.remote(config, data_csv, num_workers, outdir))
     assert actual_lr == expected_lr
+
+
+def test_changing_parameters_on_plateau(tmpdir, ray_test_cluster):
+    input_features = [sequence_feature(reduce_output="sum")]
+    output_features = [category_feature(vocab_size=2, reduce_input="sum")]
+
+    csv_filename = os.path.join(tmpdir, "training.csv")
+    data_csv = generate_data(input_features, output_features, csv_filename)
+    val_csv = shutil.copyfile(data_csv, os.path.join(tmpdir, "validation.csv"))
+    test_csv = shutil.copyfile(data_csv, os.path.join(tmpdir, "test.csv"))
+    config = {
+        "input_features": input_features,
+        "output_features": output_features,
+        "combiner": {"type": "concat", "output_size": 14},
+        TRAINER: {
+            "epochs": 2,
+            "learning_rate": 1.0,
+            "reduce_learning_rate_on_plateau": 1,
+            "increase_batch_size_on_plateau": 1,
+        },
+    }
+    model = LudwigModel(config, backend=LocalTestBackend())
+
+    model.train(training_set=data_csv, validation_set=val_csv, test_set=test_csv, output_directory=tmpdir)


Code review comments:
diff_hunk: @@ -137,3 +135,27 @@ def test_scale_lr(learning_rate_scaling, expected_lr, tmpdir, ray_test_cluster):
 
     actual_lr = ray.get(run_scale_lr.remote(config, data_csv, num_workers, outdir))
     assert actual_lr == expected_lr
+
+
+def test_changing_parameters_on_plateau(tmpdir, ray_test_cluster):
path: tests/integration_tests/test_trainer.py
text: @user1:
nit: ray_test_cluster doesn't seem used here