fix tests

Author: deependujha

src/lightning/pytorch/tuner/batch_size_scaling.py

@@ -313,6 +313,8 @@ def _run_binsearch_scaling(
                 f"Applying margin of {margin:.1%}, reducing batch size from {new_size} to {margin_reduced_size}"
             )
             new_size = margin_reduced_size
+            # update attribute in the model/datamodule as well
+            lightning_setattr(trainer.lightning_module, batch_arg_name, new_size)
 
     return new_size
 
@@ -353,7 +355,10 @@ def _adjust_batch_size(
     try:
         combined_dataset_length = combined_loader._dataset_length()
         if batch_size >= combined_dataset_length:
-            rank_zero_info(f"The batch size {batch_size} is greater or equal than the length of your dataset.")
+            rank_zero_info(
+                f"The batch size {batch_size} is greater or equal than"
+                f" the length of your dataset: {combined_dataset_length}."
+            )
             return batch_size, False
     except NotImplementedError:
         # all datasets are iterable style

tests/tests_pytorch/tuner/test_scale_batch_size.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 import glob
 import logging
+import math
 import os
 from copy import deepcopy
 from unittest.mock import patch
@@ -69,7 +70,7 @@ def test_scale_batch_size_method_with_model_or_datamodule(tmp_path, model_bs, dm
 
     tuner = Tuner(trainer)
     new_batch_size = tuner.scale_batch_size(model, mode="binsearch", init_val=4, max_trials=2, datamodule=datamodule)
-    assert new_batch_size == 8
+    assert new_batch_size == 7  # applied margin of 5% on 8 -> int(8 * 0.95) = 7
 
     if model_bs is not None:
         assert model.batch_size == new_batch_size
@@ -317,7 +318,12 @@ def test_dataloader_reset_with_scale_batch_size(tmp_path, caplog, scale_method,
     # With our fix, when max_trials is reached, we don't try the doubled batch size, so we get max_trials - 1 messages
     expected_tries = max_trials - 1 if init_batch_size < dataset_len and max_trials > 0 else 0
     assert caplog.text.count("trying batch size") == expected_tries
-    assert caplog.text.count("greater or equal than the length") == int(new_batch_size == dataset_len)
+
+    # Determine the largest batch size that was actually tested.
+    # For "power" this is the final found size; for "binsearch" we applied a 5% margin
+    # when storing the final value, so the largest tested value is the one before applying margin.
+    largest_tested_batch_size = new_batch_size if scale_method == "power" else int(math.ceil(new_batch_size * 100 / 95))
+    assert caplog.text.count("greater or equal than the length") == int(largest_tested_batch_size >= dataset_len)
 
     assert trainer.train_dataloader.batch_size == new_batch_size
     assert trainer.val_dataloaders.batch_size == new_batch_size
@@ -453,7 +459,7 @@ def val_dataloader(self):
         tuner.scale_batch_size(model, method="validate")
 
 
-@pytest.mark.parametrize(("scale_method", "expected_batch_size"), [("power", 62), ("binsearch", 100)])
+@pytest.mark.parametrize(("scale_method", "expected_batch_size"), [("power", 62), ("binsearch", 95)])
 @patch("lightning.pytorch.tuner.batch_size_scaling.is_oom_error", return_value=True)
 def test_dataloader_batch_size_updated_on_failure(_, tmp_path, scale_method, expected_batch_size):
     class CustomBatchSizeModel(BatchSizeModel):
@@ -611,7 +617,7 @@ def training_step(self, batch, batch_idx):
     ("max_trials", "mode", "init_val", "expected"),
     [
         (3, "power", 2, 8),
-        (3, "binsearch", 2, 8),
+        (3, "binsearch", 2, 7),  # applied margin of 5% on 8 -> int(8 * 0.95) = 7
         (1, "power", 4, 4),
         (0, "power", 2, 2),
     ],