From 7cf9005b287b92ec10d183a12aae9b4570adb586 Mon Sep 17 00:00:00 2001 From: Gokul Gunasekaran Date: Tue, 26 Mar 2024 14:28:27 -0700 Subject: [PATCH] Minor changes --- test/dataloader2/test_mprs.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/test/dataloader2/test_mprs.py b/test/dataloader2/test_mprs.py index 4f0878108..83a1114b6 100644 --- a/test/dataloader2/test_mprs.py +++ b/test/dataloader2/test_mprs.py @@ -229,7 +229,8 @@ class TestMultiProcessingReadingService(TestCase): @parametrize("dp_fn", [subtest(_non_dispatching_dp, "non_dispatch"), subtest(_dispatching_dp, "dispatch")]) @parametrize("main_prefetch", [0, 10]) @parametrize("worker_prefetch", [0, 10]) - def test_early_exit(self, ctx="forkserver", dp_fn, main_prefetch, worker_prefetch) -> None: + def test_early_exit(self, dp_fn, main_prefetch, worker_prefetch) -> None: + ctx = "forkserver" dp = dp_fn(1000) rs = MultiProcessingReadingService( num_workers=2, @@ -246,7 +247,8 @@ def test_early_exit(self, ctx="forkserver", dp_fn, main_prefetch, worker_prefetc @parametrize("dp_fn", [subtest(_non_dispatching_dp, "non_dispatch"), subtest(_dispatching_dp, "dispatch")]) @parametrize("main_prefetch", [0, 10]) @parametrize("worker_prefetch", [0, 10]) - def test_exit(self, ctx="forkserver", dp_fn, main_prefetch, worker_prefetch) -> None: + def test_exit(self, dp_fn, main_prefetch, worker_prefetch) -> None: + ctx = "forkserver" dp = dp_fn(1000) rs = MultiProcessingReadingService( num_workers=2, @@ -263,8 +265,8 @@ def test_exit(self, ctx="forkserver", dp_fn, main_prefetch, worker_prefetch) -> "n_workers,worker_prefetch_cnt,main_prefetch_cnt", [(1, 0, 0), (1, 0, 2), (2, 0, 0), (2, 2, 0), (2, 0, 2), (2, 2, 2)], ) - def test_reading_service_pause_resume(self, ctx="forkserver", dp, n_workers, worker_prefetch_cnt, main_prefetch_cnt) -> None: - + def test_reading_service_pause_resume(self, dp, n_workers, worker_prefetch_cnt, main_prefetch_cnt) -> None: + ctx = "forkserver" # Functional Test: Testing various configuration of DataPipe/ReadingService to ensure the pipeline # properly pauses and resumes rs = MultiProcessingReadingService( @@ -292,8 +294,8 @@ def test_reading_service_pause_resume(self, ctx="forkserver", dp, n_workers, wor @dp_parametrize @parametrize("n_workers,worker_prefetch_cnt,main_prefetch_cnt", [(2, 0, 1), (2, 1, 0), (2, 0, 0)]) - def test_reading_service_pause_stop_yield(self, ctx="forkserver", dp, n_workers, worker_prefetch_cnt, main_prefetch_cnt) -> None: - + def test_reading_service_pause_stop_yield(self, dp, n_workers, worker_prefetch_cnt, main_prefetch_cnt) -> None: + ctx = "forkserver" # Functional Test: Confirms that `dl` will stop yielding elements after `_pause` is called rs = MultiProcessingReadingService( num_workers=n_workers,