[flaky test] Fix flaky checkpointing tests (#5791)

* Fix flaky checkpointing tests

* Fix checkpoint test logic

* Fix exception matching

* timeout exception

* Fix import

* fix build

python/ray/tests/test_basic.py

@@ -33,6 +33,8 @@ import ray.ray_constants as ray_constants
 import ray.tests.cluster_utils
 import ray.tests.utils
 
+from ray.tests.utils import RayTestTimeoutException
+
 logger = logging.getLogger(__name__)
 
 
@@ -1219,8 +1221,9 @@ def test_profiling_api(ray_start_2_cpus):
     start_time = time.time()
     while True:
         if time.time() - start_time > timeout_seconds:
-            raise Exception("Timed out while waiting for information in "
-                            "profile table.")
+            raise RayTestTimeoutException(
+                "Timed out while waiting for information in "
+                "profile table.")
         profile_data = ray.timeline()
         event_types = {event["cat"] for event in profile_data}
         expected_types = [
@@ -1921,8 +1924,9 @@ def test_gpu_ids(shutdown_only):
         if len(set(ray.get([f.remote() for _ in range(10)]))) == 10:
             break
         if time.time() > start_time + 10:
-            raise Exception("Timed out while waiting for workers to start "
-                            "up.")
+            raise RayTestTimeoutException(
+                "Timed out while waiting for workers to start "
+                "up.")
 
     list_of_ids = ray.get([f0.remote() for _ in range(10)])
     assert list_of_ids == 10 * [[]]
@@ -2519,7 +2523,7 @@ def wait_for_num_tasks(num_tasks, timeout=10):
         if len(ray.tasks()) >= num_tasks:
             return
         time.sleep(0.1)
-    raise Exception("Timed out while waiting for global state.")
+    raise RayTestTimeoutException("Timed out while waiting for global state.")
 
 
 def wait_for_num_objects(num_objects, timeout=10):
@@ -2528,7 +2532,7 @@ def wait_for_num_objects(num_objects, timeout=10):
         if len(ray.objects()) >= num_objects:
             return
         time.sleep(0.1)
-    raise Exception("Timed out while waiting for global state.")
+    raise RayTestTimeoutException("Timed out while waiting for global state.")
 
 
 @pytest.mark.skipif(
@@ -2621,8 +2625,9 @@ def test_global_state_api(shutdown_only):
             if tables_ready:
                 return
             time.sleep(0.1)
-        raise Exception("Timed out while waiting for object table to "
-                        "update.")
+        raise RayTestTimeoutException(
+            "Timed out while waiting for object table to "
+            "update.")
 
     object_table = ray.objects()
     assert len(object_table) == 2