randomize order of pytests

mosaicml · Feb 9, 2023 · 4f3bff9 · 4f3bff9
1 parent a918b2c
commit 4f3bff9
Showing 1 changed file with 61 additions and 61 deletions.
diff --git a/tests/trainer/test_trainer.py b/tests/trainer/test_trainer.py
@@ -893,6 +893,67 @@ def test_training_duration_unit(
                 assert event_counter_callback.event_to_num_calls[Event.EPOCH_CHECKPOINT] == 2
 
 
+@pytest.mark.vision
+class TestFFCVDataloaders:
+
+    train_file = None
+    val_file = None
+    tmp_path = None
+
+    @pytest.fixture(autouse=True)
+    def create_dataset(self, tmp_path_factory: pytest.TempPathFactory):
+        dataset_train = RandomImageDataset(size=16, is_PIL=True)
+        self.tmp_path = tmp_path_factory.mktemp('ffcv')
+        output_train_file = str(self.tmp_path / 'train.ffcv')
+        write_ffcv_dataset(dataset_train, write_path=output_train_file, num_workers=1, write_mode='proportion')
+        dataset_val = RandomImageDataset(size=16, is_PIL=True)
+        output_val_file = str(self.tmp_path / 'val.ffcv')
+        write_ffcv_dataset(dataset_val, write_path=output_val_file, num_workers=1, write_mode='proportion')
+        self.train_file = output_train_file
+        self.val_file = output_val_file
+
+    def _get_dataloader(self, is_train):
+        assert self.tmp_path is not None
+        assert self.train_file is not None
+        assert self.val_file is not None
+        datadir = os.path.join(self.tmp_path, self.train_file if is_train else self.val_file)
+        return build_ffcv_imagenet_dataloader(
+            datadir=str(datadir),
+            global_batch_size=4,
+            is_train=is_train,
+            num_workers=0,
+        )
+
+    @pytest.fixture
+    def config(self):
+        try:
+            import ffcv
+        except ImportError as e:
+            raise ImportError(('Composer was installed without ffcv support. '
+                               'To use ffcv with Composer, please install ffcv in your environment.')) from e
+        train_dataloader = self._get_dataloader(is_train=True)
+        val_dataloader = self._get_dataloader(is_train=False)
+        assert isinstance(train_dataloader, ffcv.Loader)
+        assert isinstance(val_dataloader, ffcv.Loader)
+        return {
+            'model': SimpleConvModel(),
+            'train_dataloader': train_dataloader,
+            'eval_dataloader': val_dataloader,
+            'max_duration': '2ep',
+        }
+
+    """
+    Tests that training completes with ffcv dataloaders.
+    """
+
+    @device('gpu-amp', precision=True)
+    def test_ffcv(self, config, device, precision):
+        config['device'] = device
+        config['precision'] = precision
+        trainer = Trainer(**config)
+        trainer.fit()
+
+
 @world_size(1, 2)
 @device('cpu', 'gpu', 'gpu-amp', precision=True)
 class TestTrainerEquivalence():
@@ -1144,67 +1205,6 @@ def test_data_not_augmented(self, config):
             trainer.fit()
 
 
-@pytest.mark.vision
-class TestFFCVDataloaders:
-
-    train_file = None
-    val_file = None
-    tmp_path = None
-
-    @pytest.fixture(autouse=True)
-    def create_dataset(self, tmp_path_factory: pytest.TempPathFactory):
-        dataset_train = RandomImageDataset(size=16, is_PIL=True)
-        self.tmp_path = tmp_path_factory.mktemp('ffcv')
-        output_train_file = str(self.tmp_path / 'train.ffcv')
-        write_ffcv_dataset(dataset_train, write_path=output_train_file, num_workers=1, write_mode='proportion')
-        dataset_val = RandomImageDataset(size=16, is_PIL=True)
-        output_val_file = str(self.tmp_path / 'val.ffcv')
-        write_ffcv_dataset(dataset_val, write_path=output_val_file, num_workers=1, write_mode='proportion')
-        self.train_file = output_train_file
-        self.val_file = output_val_file
-
-    def _get_dataloader(self, is_train):
-        assert self.tmp_path is not None
-        assert self.train_file is not None
-        assert self.val_file is not None
-        datadir = os.path.join(self.tmp_path, self.train_file if is_train else self.val_file)
-        return build_ffcv_imagenet_dataloader(
-            datadir=str(datadir),
-            global_batch_size=4,
-            is_train=is_train,
-            num_workers=0,
-        )
-
-    @pytest.fixture
-    def config(self):
-        try:
-            import ffcv
-        except ImportError as e:
-            raise ImportError(('Composer was installed without ffcv support. '
-                               'To use ffcv with Composer, please install ffcv in your environment.')) from e
-        train_dataloader = self._get_dataloader(is_train=True)
-        val_dataloader = self._get_dataloader(is_train=False)
-        assert isinstance(train_dataloader, ffcv.Loader)
-        assert isinstance(val_dataloader, ffcv.Loader)
-        return {
-            'model': SimpleConvModel(),
-            'train_dataloader': train_dataloader,
-            'eval_dataloader': val_dataloader,
-            'max_duration': '2ep',
-        }
-
-    """
-    Tests that training completes with ffcv dataloaders.
-    """
-
-    @device('gpu-amp', precision=True)
-    def test_ffcv(self, config, device, precision):
-        config['device'] = device
-        config['precision'] = precision
-        trainer = Trainer(**config)
-        trainer.fit()
-
-
 @pytest.mark.world_size(2)
 def test_state_run_name():
     # seeding with the global rank to ensure that each rank has a different seed