852 adds a timeout call decorator

tests/test_highresnet.py

@@ -15,7 +15,7 @@
 from parameterized import parameterized
 
 from monai.networks.nets import HighResNet
-from tests.utils import test_script_save
+from tests.utils import DistTestCase, TimedCall, test_script_save
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -44,7 +44,7 @@
 ]
 
 
-class TestHighResNet(unittest.TestCase):
+class TestHighResNet(DistTestCase):
     @parameterized.expand([TEST_CASE_1, TEST_CASE_2, TEST_CASE_3, TEST_CASE_4])
     def test_shape(self, input_param, input_shape, expected_shape):
         net = HighResNet(**input_param).to(device)
@@ -53,6 +53,7 @@ def test_shape(self, input_param, input_shape, expected_shape):
             result = net.forward(torch.randn(input_shape).to(device))
             self.assertEqual(result.shape, expected_shape)
 
+    @TimedCall(seconds=100, force_quit=True)
     def test_script(self):
         input_param, input_shape, expected_shape = TEST_CASE_1
         net = HighResNet(**input_param)

tests/test_integration_classification_2d.py

@@ -25,7 +25,7 @@
 from monai.transforms import AddChannel, Compose, LoadPNG, RandFlip, RandRotate, RandZoom, ScaleIntensity, ToTensor
 from monai.utils import set_determinism
 from tests.testing_data.integration_answers import test_integration_value
-from tests.utils import skip_if_quick
+from tests.utils import DistTestCase, TimedCall, skip_if_quick
 
 TEST_DATA_URL = "https://www.dropbox.com/s/5wwskxctvcxiuea/MedNIST.tar.gz?dl=1"
 MD5_VALUE = "0bc7306e7427e00ad1c5526a6677552d"
@@ -45,7 +45,7 @@ def __getitem__(self, index):
         return self.transforms(self.image_files[index]), self.labels[index]
 
 
-def run_training_test(root_dir, train_x, train_y, val_x, val_y, device="cuda:0"):
+def run_training_test(root_dir, train_x, train_y, val_x, val_y, device="cuda:0", num_workers=10):
 
     monai.config.print_config()
     # define transforms for image and classification
@@ -65,10 +65,10 @@ def run_training_test(root_dir, train_x, train_y, val_x, val_y, device="cuda:0")
 
     # create train, val data loaders
     train_ds = MedNISTDataset(train_x, train_y, train_transforms)
-    train_loader = DataLoader(train_ds, batch_size=300, shuffle=True, num_workers=10)
+    train_loader = DataLoader(train_ds, batch_size=300, shuffle=True, num_workers=num_workers)
 
     val_ds = MedNISTDataset(val_x, val_y, val_transforms)
-    val_loader = DataLoader(val_ds, batch_size=300, num_workers=10)
+    val_loader = DataLoader(val_ds, batch_size=300, num_workers=num_workers)
 
     model = densenet121(spatial_dims=2, in_channels=1, out_channels=len(np.unique(train_y))).to(device)
     loss_function = torch.nn.CrossEntropyLoss()
@@ -127,11 +127,11 @@ def run_training_test(root_dir, train_x, train_y, val_x, val_y, device="cuda:0")
     return epoch_loss_values, best_metric, best_metric_epoch
 
 
-def run_inference_test(root_dir, test_x, test_y, device="cuda:0"):
+def run_inference_test(root_dir, test_x, test_y, device="cuda:0", num_workers=10):
     # define transforms for image and classification
     val_transforms = Compose([LoadPNG(image_only=True), AddChannel(), ScaleIntensity(), ToTensor()])
     val_ds = MedNISTDataset(test_x, test_y, val_transforms)
-    val_loader = DataLoader(val_ds, batch_size=300, num_workers=10)
+    val_loader = DataLoader(val_ds, batch_size=300, num_workers=num_workers)
 
     model = densenet121(spatial_dims=2, in_channels=1, out_channels=len(np.unique(test_y))).to(device)
 
@@ -152,7 +152,7 @@ def run_inference_test(root_dir, test_x, test_y, device="cuda:0"):
 
 
 @skip_if_quick
-class IntegrationClassification2D(unittest.TestCase):
+class IntegrationClassification2D(DistTestCase):
     def setUp(self):
         set_determinism(seed=0)
         self.data_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "testing_data")
@@ -198,7 +198,7 @@ def setUp(self):
                 self.train_x.append(image_file_list[i])
                 self.train_y.append(image_classes[i])
 
-        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu:0")
+        self.device = "cuda:0" if torch.cuda.is_available() else "cpu:0"
 
     def tearDown(self):
         set_determinism(seed=None)
@@ -208,38 +208,45 @@ def tearDown(self):
             warnings.warn("not found best_metric_model.pth, training skipped?")
             pass
 
-    def test_training(self):
+    def train_and_infer(self, idx=0):
+        results = []
         if not os.path.exists(os.path.join(self.data_dir, "MedNIST")):
             # skip test if no MedNIST dataset
-            return
+            return results
+
+        set_determinism(seed=0)
+        losses, best_metric, best_metric_epoch = run_training_test(
+            self.data_dir, self.train_x, self.train_y, self.val_x, self.val_y, device=self.device
+        )
+        infer_metric = run_inference_test(self.data_dir, self.test_x, self.test_y, device=self.device)
+
+        print(f"integration_classification_2d {losses}")
+        print("best metric", best_metric)
+        print("infer metric", infer_metric)
+        # check training properties
+        self.assertTrue(test_integration_value(TASK, key="losses", data=losses, rtol=1e-2))
+        self.assertTrue(test_integration_value(TASK, key="best_metric", data=best_metric, rtol=1e-4))
+        np.testing.assert_allclose(best_metric_epoch, 4)
+        model_file = os.path.join(self.data_dir, "best_metric_model.pth")
+        self.assertTrue(os.path.exists(model_file))
+        # check inference properties
+        self.assertTrue(test_integration_value(TASK, key="infer_prop", data=np.asarray(infer_metric), rtol=1))
+        results.extend(losses)
+        results.append(best_metric)
+        results.extend(infer_metric)
+        return results
+
+    def test_training(self):
         repeated = []
         for i in range(2):
-            set_determinism(seed=0)
-
-            repeated.append([])
-            losses, best_metric, best_metric_epoch = run_training_test(
-                self.data_dir, self.train_x, self.train_y, self.val_x, self.val_y, device=self.device
-            )
-
-            # check training properties
-            print(f"integration_classification_2d {losses}")
-            self.assertTrue(test_integration_value(TASK, key="losses", data=losses, rtol=1e-2))
-            repeated[i].extend(losses)
-            print("best metric", best_metric)
-            self.assertTrue(test_integration_value(TASK, key="best_metric", data=best_metric, rtol=1e-4))
-            repeated[i].append(best_metric)
-            np.testing.assert_allclose(best_metric_epoch, 4)
-            model_file = os.path.join(self.data_dir, "best_metric_model.pth")
-            self.assertTrue(os.path.exists(model_file))
-
-            infer_metric = run_inference_test(self.data_dir, self.test_x, self.test_y, device=self.device)
-            print("infer metric", infer_metric)
-            # check inference properties
-            self.assertTrue(test_integration_value(TASK, key="infer_prop", data=np.asarray(infer_metric), rtol=1))
-            repeated[i].extend(infer_metric)
-
+            results = self.train_and_infer(i)
+            repeated.append(results)
         np.testing.assert_allclose(repeated[0], repeated[1])
 
+    @TimedCall(seconds=500, skip_timing=not torch.cuda.is_available(), daemon=False)
+    def test_timing(self):
+        self.train_and_infer()
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_integration_determinism.py

@@ -20,6 +20,7 @@
 from monai.networks.nets import UNet
 from monai.transforms import AddChannel, Compose, RandRotate90, RandSpatialCrop, ScaleIntensity, ToTensor
 from monai.utils import set_determinism
+from tests.utils import DistTestCase, TimedCall
 
 
 def run_test(batch_size=64, train_steps=200, device="cuda:0"):
@@ -67,15 +68,16 @@ def __len__(self):
     return epoch_loss, step
 
 
-class TestDeterminism(unittest.TestCase):
+class TestDeterminism(DistTestCase):
     def setUp(self):
-        set_determinism(seed=0)
         self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu:0")
 
     def tearDown(self):
         set_determinism(seed=None)
 
+    @TimedCall(seconds=30)
     def test_training(self):
+        set_determinism(seed=0)
         loss, step = run_test(device=self.device)
         print(f"Deterministic loss {loss} at training step {step}")
         np.testing.assert_allclose(step, 4)

tests/test_integration_segmentation_3d.py

@@ -40,7 +40,7 @@
 from monai.utils import set_determinism
 from monai.visualize import plot_2d_or_3d_image
 from tests.testing_data.integration_answers import test_integration_value
-from tests.utils import skip_if_quick
+from tests.utils import DistTestCase, TimedCall, skip_if_quick
 
 TASK = "integration_segmentation_3d"
 
@@ -191,7 +191,7 @@ def run_inference_test(root_dir, device="cuda:0"):
         ]
     )
     val_ds = monai.data.Dataset(data=val_files, transform=val_transforms)
-    # sliding window inferene need to input 1 image in every iteration
+    # sliding window inference need to input 1 image in every iteration
     val_loader = monai.data.DataLoader(val_ds, batch_size=1, num_workers=4)
     val_post_tran = Compose([Activations(sigmoid=True), AsDiscrete(threshold_values=True)])
     dice_metric = DiceMetric(include_background=True, reduction="mean")
@@ -228,7 +228,7 @@ def run_inference_test(root_dir, device="cuda:0"):
 
 
 @skip_if_quick
-class IntegrationSegmentation3D(unittest.TestCase):
+class IntegrationSegmentation3D(DistTestCase):
     def setUp(self):
         set_determinism(seed=0)
 
@@ -240,47 +240,53 @@ def setUp(self):
             n = nib.Nifti1Image(seg, np.eye(4))
             nib.save(n, os.path.join(self.data_dir, f"seg{i:d}.nii.gz"))
 
-        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu:0")
+        self.device = "cuda:0" if torch.cuda.is_available() else "cpu:0"
 
     def tearDown(self):
         set_determinism(seed=None)
         shutil.rmtree(self.data_dir)
 
-    def test_training(self):
-        repeated = []
-        for i in range(4):
-            set_determinism(0)
-
-            repeated.append([])
-            losses, best_metric, best_metric_epoch = run_training_test(
-                self.data_dir, device=self.device, cachedataset=i
-            )
+    def train_and_infer(self, idx=0):
+        results = []
+        set_determinism(0)
+        losses, best_metric, best_metric_epoch = run_training_test(self.data_dir, device=self.device, cachedataset=idx)
+        infer_metric = run_inference_test(self.data_dir, device=self.device)
 
-            # check training properties
-            print("losses", losses)
-            self.assertTrue(test_integration_value(TASK, key="losses", data=losses, rtol=1e-3))
-            repeated[i].extend(losses)
-            print("best metric", best_metric)
-            self.assertTrue(test_integration_value(TASK, key="best_metric", data=best_metric, rtol=1e-2))
-            repeated[i].append(best_metric)
-            self.assertTrue(len(glob(os.path.join(self.data_dir, "runs"))) > 0)
-            model_file = os.path.join(self.data_dir, "best_metric_model.pth")
-            self.assertTrue(os.path.exists(model_file))
+        # check training properties
+        print("losses", losses)
+        print("best metric", best_metric)
+        print("infer metric", infer_metric)
+        self.assertTrue(test_integration_value(TASK, key="losses", data=losses, rtol=1e-3))
+        self.assertTrue(test_integration_value(TASK, key="best_metric", data=best_metric, rtol=1e-2))
+        self.assertTrue(len(glob(os.path.join(self.data_dir, "runs"))) > 0)
+        model_file = os.path.join(self.data_dir, "best_metric_model.pth")
+        self.assertTrue(os.path.exists(model_file))
 
-            infer_metric = run_inference_test(self.data_dir, device=self.device)
+        # check inference properties
+        self.assertTrue(test_integration_value(TASK, key="infer_metric", data=infer_metric, rtol=1e-2))
+        output_files = sorted(glob(os.path.join(self.data_dir, "output", "img*", "*.nii.gz")))
+        print([np.mean(nib.load(output).get_fdata()) for output in output_files])
+        results.extend(losses)
+        results.append(best_metric)
+        results.append(infer_metric)
+        for output in output_files:
+            ave = np.mean(nib.load(output).get_fdata())
+            results.append(ave)
+        self.assertTrue(test_integration_value(TASK, key="output_sums", data=results[8:], rtol=1e-2))
+        return results
 
-            # check inference properties
-            print("infer metric", infer_metric)
-            self.assertTrue(test_integration_value(TASK, key="infer_metric", data=infer_metric, rtol=1e-2))
-            repeated[i].append(infer_metric)
-            output_files = sorted(glob(os.path.join(self.data_dir, "output", "img*", "*.nii.gz")))
-            print([np.mean(nib.load(output).get_fdata()) for output in output_files])
-            for output in output_files:
-                ave = np.mean(nib.load(output).get_fdata())
-                repeated[i].append(ave)
-            self.assertTrue(test_integration_value(TASK, key="output_sums", data=repeated[i][8:], rtol=1e-2))
+    def test_training(self):
+        repeated = []
+        for i in range(4):
+            results = self.train_and_infer(i)
+            repeated.append(results)
         np.testing.assert_allclose(repeated[0], repeated[1])
         np.testing.assert_allclose(repeated[0], repeated[2])
+        np.testing.assert_allclose(repeated[0], repeated[3])
+
+    @TimedCall(seconds=180, daemon=False)
+    def test_timing(self):
+        self.train_and_infer(idx=3)
 
 
 if __name__ == "__main__":

tests/test_integration_sliding_window.py

@@ -26,7 +26,7 @@
 from monai.networks.nets import UNet
 from monai.transforms import AddChannel
 from monai.utils import set_determinism
-from tests.utils import make_nifti_image, skip_if_quick
+from tests.utils import DistTestCase, TimedCall, make_nifti_image, skip_if_quick
 
 
 def run_test(batch_size, img_name, seg_name, output_dir, device="cuda:0"):
@@ -60,7 +60,7 @@ def _sliding_window_processor(_engine, batch):
 
 
 @skip_if_quick
-class TestIntegrationSlidingWindow(unittest.TestCase):
+class TestIntegrationSlidingWindow(DistTestCase):
     def setUp(self):
         set_determinism(seed=0)
 
@@ -76,7 +76,9 @@ def tearDown(self):
         if os.path.exists(self.seg_name):
             os.remove(self.seg_name)
 
+    @TimedCall(seconds=10)
     def test_training(self):
+        set_determinism(seed=0)
         with tempfile.TemporaryDirectory() as tempdir:
             output_file = run_test(
                 batch_size=2, img_name=self.img_name, seg_name=self.seg_name, output_dir=tempdir, device=self.device

tests/test_integration_stn.py

@@ -22,6 +22,7 @@
 from monai.data import create_test_image_2d
 from monai.networks.layers import AffineTransform
 from monai.utils import set_determinism
+from tests.utils import DistTestCase, TimedCall
 
 
 class STNBenchmark(nn.Module):
@@ -96,13 +97,14 @@ def compare_2d(is_ref=True, device=None, reverse_indexing=False):
     return model(img_a).detach().cpu().numpy(), loss.item(), init_loss
 
 
-class TestSpatialTransformerCore(unittest.TestCase):
+class TestSpatialTransformerCore(DistTestCase):
     def setUp(self):
         self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu:0")
 
     def tearDown(self):
         set_determinism(seed=None)
 
+    @TimedCall(seconds=60)
     def test_training(self):
         """
         check that the quality AffineTransform backpropagation

tests/test_integration_unet_2d.py

@@ -19,7 +19,7 @@
 from monai.data import create_test_image_2d
 from monai.losses import DiceLoss
 from monai.networks.nets import BasicUNet, UNet
-from tests.utils import skip_if_quick
+from tests.utils import DistTestCase, TimedCall, skip_if_quick
 
 
 def run_test(net_name="basicunet", batch_size=64, train_steps=100, device="cuda:0"):
@@ -51,7 +51,8 @@ def __len__(self):
 
 
 @skip_if_quick
-class TestIntegrationUnet2D(unittest.TestCase):
+class TestIntegrationUnet2D(DistTestCase):
+    @TimedCall(seconds=20, daemon=False)
     def test_unet_training(self):
         for n in ["basicunet", "unet"]:
             loss = run_test(net_name=n, device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu:0"))