Refinements to the num-workers warning (#18737)

Author: awaelchli

src/lightning/fabric/utilities/data.py

@@ -447,7 +447,7 @@ def suggested_max_num_workers(local_world_size: int) -> int:
     if local_world_size < 1:
         raise ValueError(f"`local_world_size` should be >= 1, got {local_world_size}.")
     cpu_count = _num_cpus_available()
-    return max(1, cpu_count // local_world_size)
+    return max(1, cpu_count // local_world_size - 1)  # -1 to leave some resources for main process
 
 
 def _num_cpus_available() -> int:

src/lightning/pytorch/trainer/connectors/data_connector.py

@@ -436,7 +436,7 @@ def _worker_check(trainer: "pl.Trainer", dataloader: object, name: str) -> None:
         rank_zero_warn(
             f"Consider setting `persistent_workers=True` in '{name}' to speed up the dataloader worker initialization."
         )
-    elif dataloader.num_workers <= 2 < upper_bound or dataloader.num_workers < 2 <= upper_bound:
+    elif dataloader.num_workers < 2:
         # if changed, update the `filterwarnings` snippet in 'speed.html#num-workers'
         rank_zero_warn(
             f"The '{name}' does not have many workers which may be a bottleneck. Consider increasing the value of the"

tests/tests_fabric/utilities/test_data.py

@@ -588,13 +588,13 @@ def test_set_sampler_epoch():
     [
         (0, 1, 1),
         (1, 1, 1),
-        (2, 1, 2),
+        (2, 1, 2 - 1),
         (1, 2, 1),
         (2, 2, 1),
         (3, 2, 1),
-        (4, 2, 2),
+        (4, 2, 2 - 1),
         (4, 3, 1),
-        (4, 1, 4),
+        (4, 1, 4 - 1),
     ],
 )
 @pytest.mark.parametrize(

tests/tests_pytorch/trainer/connectors/test_data_connector.py

@@ -139,31 +139,16 @@ def test_dataloader_persistent_workers_performance_warning(num_workers, tmp_path
     trainer.fit(model, dataloader)
 
 
-@pytest.mark.parametrize(
-    ("num_devices", "num_workers", "cpu_count", "expected_warning"),
-    [
-        (1, 0, 1, False),
-        (8, 0, 1, False),
-        (8, 0, None, False),
-        (1, 1, None, False),
-        (1, 2, 2, False),
-        (1, 1, 8, True),
-        (1, 2, 8, True),
-        (1, 3, 8, False),
-        (4, 1, 8, True),
-        (4, 2, 8, False),
-        (8, 2, 8, False),
-    ],
-)
+@pytest.mark.parametrize(("num_workers", "expected_warning"), [(0, True), (1, True), (2, False), (3, False)])
 @mock.patch("lightning.fabric.utilities.data.os.cpu_count")
 @mock.patch("lightning.pytorch.trainer.connectors.data_connector.mp.get_start_method", return_value="not_spawn")
-def test_worker_check(_, cpu_count_mock, num_devices, num_workers, cpu_count, expected_warning, monkeypatch):
+def test_worker_check(_, cpu_count_mock, num_workers, expected_warning, monkeypatch):
     monkeypatch.delattr(lightning.fabric.utilities.data.os, "sched_getaffinity", raising=False)
     trainer = Mock(spec=Trainer)
     dataloader = Mock(spec=DataLoader, persistent_workers=False)
-    trainer.num_devices = num_devices
+    trainer.num_devices = 2
     dataloader.num_workers = num_workers
-    cpu_count_mock.return_value = cpu_count
+    cpu_count_mock.return_value = 8
 
     if expected_warning:
         ctx = pytest.warns(UserWarning, match="Consider increasing the value of the `num_workers` argument`")