update combined test

Lightning-AI · Jul 17, 2024 · e621185 · e621185
1 parent f52a501
commit e621185
Showing 1 changed file with 8 additions and 82 deletions.
diff --git a/tests/streaming/test_combined.py b/tests/streaming/test_combined.py
@@ -568,43 +568,6 @@ def test_combined_dataset_with_dataloader_2_epochs(tmpdir):
             "latest_worker_idx": 1,
             "num_samples_yielded": {0: [3, 1], 1: [3, 1], 2: [2, 0]},
         },
-        {
-            "dataset": {
-                "0": {
-                    "num_samples_yielded": 8,
-                    "num_workers": 3,
-                    "batch_size": 2,
-                    "current_epoch": 1,
-                    "input_dir_path": ANY,
-                    "input_dir_url": ANY,
-                    "item_loader": None,
-                    "drop_last": False,
-                    "seed": 42,
-                    "world_size": 1,
-                    "shuffle": True,
-                    "subsampled_files": ANY,
-                    "region_of_interest": ANY,
-                },
-                "1": {
-                    "num_samples_yielded": 3,
-                    "num_workers": 3,
-                    "batch_size": 2,
-                    "current_epoch": 1,
-                    "input_dir_path": ANY,
-                    "input_dir_url": ANY,
-                    "item_loader": None,
-                    "drop_last": False,
-                    "seed": 42,
-                    "world_size": 1,
-                    "shuffle": True,
-                    "subsampled_files": ANY,
-                    "region_of_interest": ANY,
-                },
-            },
-            "current_epoch": 0,
-            "latest_worker_idx": 2,
-            "num_samples_yielded": {0: [3, 1], 1: [3, 1], 2: [2, 1]},
-        },
         {
             "dataset": {
                 "0": {
@@ -639,8 +602,8 @@ def test_combined_dataset_with_dataloader_2_epochs(tmpdir):
                 },
             },
             "current_epoch": 0,
-            "latest_worker_idx": 0,
-            "num_samples_yielded": {0: [4, 1], 1: [3, 1], 2: [2, 1]},
+            "latest_worker_idx": 2,
+            "num_samples_yielded": {0: [3, 1], 1: [3, 1], 2: [3, 1]},
         },
         {
             "dataset": {
@@ -676,8 +639,8 @@ def test_combined_dataset_with_dataloader_2_epochs(tmpdir):
                 },
             },
             "current_epoch": 0,
-            "latest_worker_idx": 1,
-            "num_samples_yielded": {0: [4, 1], 1: [4, 1], 2: [2, 1]},
+            "latest_worker_idx": 0,
+            "num_samples_yielded": {0: [4, 1], 1: [3, 1], 2: [3, 1]},
         },
     ]
 
@@ -867,43 +830,6 @@ def test_combined_dataset_with_dataloader_2_epochs(tmpdir):
             "latest_worker_idx": 1,
             "num_samples_yielded": {0: [3, 1], 1: [3, 1], 2: [2, 0]},
         },
-        {
-            "dataset": {
-                "0": {
-                    "num_samples_yielded": 8,
-                    "num_workers": 3,
-                    "batch_size": 2,
-                    "current_epoch": 2,
-                    "input_dir_path": ANY,
-                    "input_dir_url": ANY,
-                    "item_loader": None,
-                    "drop_last": False,
-                    "seed": 42,
-                    "world_size": 1,
-                    "shuffle": True,
-                    "subsampled_files": ANY,
-                    "region_of_interest": ANY,
-                },
-                "1": {
-                    "num_samples_yielded": 3,
-                    "num_workers": 3,
-                    "batch_size": 2,
-                    "current_epoch": 2,
-                    "input_dir_path": ANY,
-                    "input_dir_url": ANY,
-                    "item_loader": None,
-                    "drop_last": False,
-                    "seed": 42,
-                    "world_size": 1,
-                    "shuffle": True,
-                    "subsampled_files": ANY,
-                    "region_of_interest": ANY,
-                },
-            },
-            "current_epoch": 1,
-            "latest_worker_idx": 2,
-            "num_samples_yielded": {0: [3, 1], 1: [3, 1], 2: [2, 1]},
-        },
         {
             "dataset": {
                 "0": {
@@ -938,8 +864,8 @@ def test_combined_dataset_with_dataloader_2_epochs(tmpdir):
                 },
             },
             "current_epoch": 1,
-            "latest_worker_idx": 0,
-            "num_samples_yielded": {0: [4, 1], 1: [3, 1], 2: [2, 1]},
+            "latest_worker_idx": 2,
+            "num_samples_yielded": {0: [3, 1], 1: [3, 1], 2: [3, 1]},
         },
         {
             "dataset": {
@@ -975,8 +901,8 @@ def test_combined_dataset_with_dataloader_2_epochs(tmpdir):
                 },
             },
             "current_epoch": 1,
-            "latest_worker_idx": 1,
-            "num_samples_yielded": {0: [4, 1], 1: [4, 1], 2: [2, 1]},
+            "latest_worker_idx": 0,
+            "num_samples_yielded": {0: [4, 1], 1: [3, 1], 2: [3, 1]},
         },
     ]