Repository: Project-MONAI/MONAI

Issue:
Revert the workaround of base image downgrading (PR 3949)
**Is your feature request related to a problem? Please describe.**
As described and implemented in https://github.com/Project-MONAI/MONAI/pull/3949, pytorch docker images after v21.10 seems to have a much larger memory footprint when launching the dataloader with num_workers>1 for the integration test:  
https://github.com/Project-MONAI/MONAI/blob/6ff85350d8cd698507e4476997e32e09b3be5eed/tests/test_integration_classification_2d.py#L59
the base image was downgraded as a temp workaround to avoid this issue in thg 16GB v100 CI.



Discussion:
the issue is with this checker that creates additional cuda contexts for each dataloader worker... https://github.com/Project-MONAI/MONAI/blob/459b1885de2c3415bc784fbb0382a4165aad63b4/tests/utils.py#L248-L266 I'll try to make a patch to avoid it.

Resolving patch:
diff --git a/monai/data/image_reader.py b/monai/data/image_reader.py
--- a/monai/data/image_reader.py
+++ b/monai/data/image_reader.py
@@ -678,7 +678,7 @@ def get_data(self, img):
         It computes `spatial_shape` and stores it in meta dict.
         When loading a list of files, they are stacked together at a new dimension as the first dimension,
         and the meta data of the first image is used to represent the output meta data.
-        Note that it will switch axis 0 and 1 after loading the array because the `HW` definition in PIL
+        Note that it will swap axis 0 and 1 after loading the array because the `HW` definition in PIL
         is different from other common medical packages.
 
         Args:


Test patch:
diff --git a/tests/test_integration_workers.py b/tests/test_integration_workers.py
new file mode 100644
--- /dev/null
+++ b/tests/test_integration_workers.py
@@ -0,0 +1,53 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+import torch
+
+from monai.data import DataLoader
+from monai.utils import set_determinism
+from tests.utils import DistTestCase, TimedCall, skip_if_no_cuda, skip_if_quick
+
+
+def run_loading_test(num_workers=50, device="cuda:0" if torch.cuda.is_available() else "cpu", pw=False):
+    """multi workers stress tests"""
+    set_determinism(seed=0)
+    train_ds = list(range(10000))
+    train_loader = DataLoader(train_ds, batch_size=300, shuffle=True, num_workers=num_workers, persistent_workers=pw)
+    answer = []
+    for _ in range(2):
+        np.testing.assert_equal(torch.cuda.memory_allocated(), 0)
+        for batch_data in train_loader:
+            x = batch_data.to(device)
+            mem = torch.cuda.memory_allocated()
+            np.testing.assert_equal(mem > 0 and mem < 5000, True)
+        answer.append(x[-1].item())
+        del x
+    return answer
+
+
+@skip_if_quick
+@skip_if_no_cuda
+class IntegrationLoading(DistTestCase):
+    def tearDown(self):
+        set_determinism(seed=None)
+
+    @TimedCall(seconds=5000, skip_timing=not torch.cuda.is_available(), daemon=False)
+    def test_timing(self):
+        for pw, expected in zip((False, True), ((6966, 7714), (6966, 4112))):
+            result = run_loading_test(pw=pw)
+            np.testing.assert_allclose(result, expected)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/utils.py b/tests/utils.py
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -245,11 +245,20 @@ def __call__(self, obj):
         )(obj)
 
 
+def is_main_test_process():
+    ps = torch.multiprocessing.current_process()
+    if not ps or not hasattr(ps, "name"):
+        return False
+    return ps.name.startswith("Main")
+
+
 def has_cupy():
     """
     Returns True if the user has installed a version of cupy.
     """
     cp, has_cp = optional_import("cupy")
+    if not is_main_test_process():
+        return has_cp  # skip the check if we are running in subprocess
     if not has_cp:
         return False
     try:  # test cupy installation with a basic example
@@ -258,7 +267,10 @@ def has_cupy():
         kernel = cp.ElementwiseKernel(
             "float32 x, float32 y", "float32 z", """ if (x - 2 > y) { z = x * y; } else { z = x + y; } """, "my_kernel"
         )
-        return kernel(x, y)[0, 0] == 0
+        flag = kernel(x, y)[0, 0] == 0
+        del x, y, kernel
+        cp.get_default_memory_pool().free_all_blocks()
+        return flag
     except Exception:
         return False