diff --git a/tests/test_ote_api.py b/tests/test_ote_api.py index 77c19342405..43ddc3e649f 100644 --- a/tests/test_ote_api.py +++ b/tests/test_ote_api.py @@ -14,6 +14,7 @@ import glob import io +import os import os.path as osp import random import time @@ -222,7 +223,11 @@ def progress_callback(progress: float, score: Optional[float] = None): # stopping process has to happen in less than 35 seconds train_future.result() self.assertEqual(training_progress_curve[-1], 100) - self.assertLess(time.time() - start_time, 100, 'Expected to stop within 100 seconds.') + if os.getenv("CUDA_VISIBLE_DEVICES") == "": + stop_timeout = 400 + else: + stop_timeout = 100 + self.assertLess(time.time() - start_time, stop_timeout, f'Expected to stop within {stop_timeout} seconds.') # Test stopping immediately (as soon as training is started). start_time = time.time() @@ -232,7 +237,12 @@ def progress_callback(progress: float, score: Optional[float] = None): detection_task.cancel_training() train_future.result() - self.assertLess(time.time() - start_time, 25) # stopping process has to happen in less than 25 seconds + if os.getenv("CUDA_VISIBLE_DEVICES") == "": + stop_process_timeout = 400 + else: + stop_process_timeout = 25 + self.assertLess(time.time() - start_time, stop_process_timeout, f'Expected to stop within {stop_process_timeout} seconds.') + @e2e_pytest_api def test_training_progress_tracking(self):