fix worker warning (Lightning-AI#2504)

* fix worker warning

* improve tests

* suggestion

Co-authored-by: Jirka <jirka@pytorchlightning.ai>

Author: jeremyjordan

pytorch_lightning/trainer/data_loading.py

@@ -124,22 +124,24 @@ def _worker_check(self, dataloader: DataLoader, name: str) -> None:
         # ddp_spawn + num_workers > 0 don't mix! tell the user
         is_dataloader = isinstance(dataloader, DataLoader)
         using_spawn = self.distributed_backend == 'ddp_spawn'
-        if is_dataloader and dataloader.num_workers > 0 and not on_windows and using_spawn:
-            rank_zero_warn('Dataloader(num_workers>0) and ddp_spawn do not mix well! '
-                           'Your performance might suffer dramatically. '
-                           'Please consider setting distributed_backend=ddp to use num_workers > 0 '
-                           '(this is a bottleneck of Python .spawn() and PyTorch')
-
-        elif is_dataloader and dataloader.num_workers <= 2 and not on_windows and not using_spawn:
-            num_cpus = multiprocessing.cpu_count()
-            rank_zero_warn(f'The dataloader, {name}, does not have many workers which may be a bottleneck.'
-                           ' Consider increasing the value of the `num_workers` argument` '
-                           f'(try {num_cpus} which is the number of cpus on this machine)'
-                           ' in the `DataLoader` init to improve performance.')
-
-        elif is_dataloader and dataloader.num_workers == 0 and not on_windows and using_spawn:
-            rank_zero_warn('You are using `distributed_backend=ddp_spawn` with num_workers=0. '
-                           'For much faster performance, switch to `distributed_backend=ddp` and set `num_workers>0`')
+        if is_dataloader and not on_windows:
+            if dataloader.num_workers > 0 and using_spawn:
+                rank_zero_warn('Dataloader(num_workers>0) and ddp_spawn do not mix well!'
+                               ' Your performance might suffer dramatically.'
+                               ' Please consider setting distributed_backend=ddp to use num_workers > 0'
+                               ' (this is a bottleneck of Python .spawn() and PyTorch')
+
+            elif dataloader.num_workers == 0 and using_spawn:
+                rank_zero_warn('You are using `distributed_backend=ddp_spawn` with num_workers=0.'
+                               ' For much faster performance, switch to `distributed_backend=ddp`'
+                               ' and set `num_workers>0`')
+
+            elif dataloader.num_workers <= 2 and multiprocessing.cpu_count() > 2 and not using_spawn:
+                num_cpus = multiprocessing.cpu_count()
+                rank_zero_warn(f'The dataloader, {name}, does not have many workers which may be a bottleneck.'
+                               ' Consider increasing the value of the `num_workers` argument`'
+                               f' (try {num_cpus} which is the number of cpus on this machine)'
+                               ' in the `DataLoader` init to improve performance.')
 
     def auto_add_sampler(self, dataloader: DataLoader, train: bool) -> DataLoader:
 

tests/trainer/test_dataloaders.py

@@ -1,14 +1,15 @@
 import platform
+from unittest.mock import patch
 
 import pytest
 import torch
 from packaging.version import parse
 from torch.utils.data.dataloader import DataLoader
-from torch.utils.data.dataset import Subset, IterableDataset
+from torch.utils.data.dataset import IterableDataset, Subset
 
 import tests.base.develop_pipelines as tpipes
 from pytorch_lightning import Trainer
-from pytorch_lightning.trainer.data_loading import _has_len, _has_iterable_dataset
+from pytorch_lightning.trainer.data_loading import _has_iterable_dataset, _has_len
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from tests.base import EvalModelTemplate
 
@@ -449,7 +450,8 @@ def test_error_on_zero_len_dataloader(tmpdir):
 
 @pytest.mark.skipif(platform.system() == 'Windows', reason='Does not apply to Windows platform.')
 @pytest.mark.parametrize('ckpt_path', [None, 'best', 'specific'])
-def test_warning_with_few_workers(tmpdir, ckpt_path):
+@patch('pytorch_lightning.trainer.data_loading.multiprocessing.cpu_count', return_value=4)
+def test_warning_with_few_workers(mock, tmpdir, ckpt_path):
     """ Test that error is raised if dataloader with only a few workers is used """
 
     model = EvalModelTemplate()
@@ -476,16 +478,22 @@ def test_warning_with_few_workers(tmpdir, ckpt_path):
     trainer = Trainer(**trainer_options)
 
     # fit model
-    with pytest.warns(UserWarning, match='train'):
+    with pytest.warns(
+        UserWarning, match='The dataloader, train dataloader, does not have many workers which may be a bottleneck.'
+    ):
         trainer.fit(model, **fit_options)
 
-    with pytest.warns(UserWarning, match='val'):
+    with pytest.warns(
+        UserWarning, match='The dataloader, val dataloader 0, does not have many workers which may be a bottleneck.'
+    ):
         trainer.fit(model, **fit_options)
 
     if ckpt_path == 'specific':
         ckpt_path = trainer.checkpoint_callback.best_model_path
     test_options = dict(test_dataloaders=train_dl, ckpt_path=ckpt_path)
-    with pytest.warns(UserWarning, match='test'):
+    with pytest.warns(
+        UserWarning, match='The dataloader, test dataloader 0, does not have many workers which may be a bottleneck.'
+    ):
         trainer.test(**test_options)