Create multiprocessing_pools once per all tests. (#1218)

Creation and destruction of multiprocessing_pools take a lot of time. So we need create them once and reuse inside all tests.

tests/conftest.py

@@ -1,11 +1,10 @@
+import multiprocessing
 import sys
 import warnings
-import multiprocessing
 
 import numpy as np
 import pytest
 
-
 try:
     import torch  # skipcq: PYL-W0611
     import torchvision  # skipcq: PYL-W0611
@@ -88,13 +87,13 @@ def float_template():
     return np.random.uniform(low=0.0, high=1.0, size=(100, 100, 3)).astype("float32")
 
 
-@pytest.fixture
-def multiprocessing_context():
+@pytest.fixture(scope="package")
+def mp_pool():
     # Usage of `fork` as a start method for multiprocessing could lead to deadlocks on macOS.
     # Because `fork` was the default start method for macOS until Python 3.8
     # we had to manually set the start method to `spawn` to avoid those issues.
     if sys.platform == "darwin":
         method = "spawn"
     else:
         method = None
-    return multiprocessing.get_context(method)
+    return multiprocessing.get_context(method).Pool(8)

tests/test_imgaug.py

@@ -3,26 +3,25 @@
 import numpy as np
 import pytest
 
+import albumentations as A
 from albumentations import Compose
 from albumentations.augmentations.bbox_utils import (
     convert_bboxes_from_albumentations,
     convert_bboxes_to_albumentations,
 )
-import albumentations as A
 from albumentations.imgaug.transforms import (
-    IAAPiecewiseAffine,
-    IAAFliplr,
-    IAAFlipud,
-    IAASuperpixels,
-    IAASharpen,
     IAAAdditiveGaussianNoise,
-    IAAPerspective,
     IAAAffine,
     IAACropAndPad,
+    IAAFliplr,
+    IAAFlipud,
+    IAAPerspective,
+    IAAPiecewiseAffine,
+    IAASharpen,
+    IAASuperpixels,
 )
 from tests.utils import set_seed
 
-
 TEST_SEEDS = (0, 1, 42, 111, 9999)
 
 
@@ -227,15 +226,12 @@ def __test_multiprocessing_support_proc(args):
         [IAAPerspective, {}],
     ],
 )
-def test_imgaug_transforms_multiprocessing_support(augmentation_cls, params, multiprocessing_context):
+def test_imgaug_transforms_multiprocessing_support(augmentation_cls, params, mp_pool):
     """Checks whether we can use augmentations in multiprocessing environments"""
     aug = augmentation_cls(p=1, **params)
     image = np.random.randint(low=0, high=256, size=(100, 100, 3), dtype=np.uint8)
 
-    pool = multiprocessing_context.Pool(8)
-    pool.map(__test_multiprocessing_support_proc, map(lambda x: (x, aug), [image] * 100))
-    pool.close()
-    pool.join()
+    mp_pool.map(__test_multiprocessing_support_proc, map(lambda x: (x, aug), [image] * 100))
 
 
 @pytest.mark.parametrize(

tests/test_random.py

@@ -1,5 +1,4 @@
 import random
-from multiprocessing.pool import Pool
 
 import numpy as np
 import pytest
@@ -25,16 +24,15 @@ def _calc(args):
         [random_utils.choice, [np.arange(1000), 100]],
     ],
 )
-def test_multiprocessing(func, args):
+def test_multiprocessing(func, args, mp_pool):
     seed = 0
     random.seed(seed)
     np.random.seed(seed)
 
     n = 10
     status = False
     for _ in range(n):
-        with Pool(2) as pool:
-            res = pool.map(_calc, [(func, args), (func, args)])
+        res = mp_pool.map(_calc, [(func, args), (func, args)])
         status = not np.allclose(res[0], res[1])
         if status:
             break

tests/test_transforms.py

@@ -228,15 +228,12 @@ def __test_multiprocessing_support_proc(args):
         },
     ),
 )
-def test_multiprocessing_support(augmentation_cls, params, multiprocessing_context):
+def test_multiprocessing_support(mp_pool, augmentation_cls, params):
     """Checks whether we can use augmentations in multiprocessing environments"""
     aug = augmentation_cls(p=1, **params)
     image = np.random.randint(low=0, high=256, size=(100, 100, 3), dtype=np.uint8)
 
-    pool = multiprocessing_context.Pool(8)
-    pool.map(__test_multiprocessing_support_proc, map(lambda x: (x, aug), [image] * 100))
-    pool.close()
-    pool.join()
+    mp_pool.map(__test_multiprocessing_support_proc, map(lambda x: (x, aug), [image] * 10))
 
 
 def test_force_apply():