Move test_timer GPU test to a separate file (#771)

Summary: Pull Request resolved: #771

Reviewed By: JKSenthil

Differential Revision: D55490411

fbshipit-source-id: e836033997ed30da4461eb530ac21993bb2d3022

Author: diego-urgell

tests/utils/test_timer.py

@@ -13,13 +13,11 @@
 from datetime import timedelta
 from random import random
 from unittest import mock
-from unittest.mock import Mock, patch
 
-import torch
 import torch.distributed as dist
 from pyre_extensions import none_throws
 from torchtnt.utils.distributed import spawn_multi_process
-from torchtnt.utils.test_utils import skip_if_not_distributed, skip_if_not_gpu
+from torchtnt.utils.test_utils import skip_if_not_distributed
 from torchtnt.utils.timer import (
     BoundedTimer,
     FullSyncPeriodicTimer,
@@ -68,8 +66,7 @@ def test_timer_context_manager_size_bound(self) -> None:
                 UPPER_BOUND,
             )
 
-    @patch("torch.cuda.synchronize")
-    def test_timer_context_manager(self, _) -> None:
+    def test_timer_context_manager(self) -> None:
         """Test the context manager in the timer class"""
 
         # Generate 3 intervals between 0.5 and 2 seconds
@@ -103,24 +100,6 @@ def test_timer_context_manager(self, _) -> None:
             timer.recorded_durations["action_4"][0], intervals[2]
         )
 
-    @skip_if_not_gpu
-    @patch("torch.cuda.synchronize")
-    def test_timer_synchronize(self, mock_synchornize: Mock) -> None:
-        """Make sure that torch.cuda.synchronize() is called when GPU is present."""
-
-        start_event = torch.cuda.Event(enable_timing=True)
-        end_event = torch.cuda.Event(enable_timing=True)
-        timer = Timer()
-
-        # Do not explicitly call synchronize, timer must call it for test to pass.
-
-        with timer.time("action_1"):
-            start_event.record()
-            time.sleep(0.5)
-            end_event.record()
-
-        self.assertEqual(mock_synchornize.call_count, 2)
-
     def test_get_timer_summary(self) -> None:
         """Test the get_timer_summary function"""
 

tests/utils/test_timer_gpu.py

@@ -0,0 +1,36 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+import time
+import unittest
+from unittest.mock import MagicMock, patch
+
+import torch
+from torchtnt.utils.test_utils import skip_if_not_gpu
+from torchtnt.utils.timer import Timer
+
+
+class TimerGPUTest(unittest.TestCase):
+    @skip_if_not_gpu
+    @patch("torch.cuda.synchronize")
+    def test_timer_synchronize(self, mock_synchronize: MagicMock) -> None:
+        """Make sure that torch.cuda.synchronize() is called when GPU is present."""
+
+        start_event = torch.cuda.Event(enable_timing=True)
+        end_event = torch.cuda.Event(enable_timing=True)
+        timer = Timer()
+
+        # Do not explicitly call synchronize, timer must call it for test to pass.
+
+        with timer.time("action_1"):
+            start_event.record()
+            time.sleep(0.5)
+            end_event.record()
+
+        self.assertEqual(mock_synchronize.call_count, 2)