[core] improve assertion check in test_task_metrics (ray-project#53958)

I'm investigating the flakiness of test_task_metrics on Windows.
Currently, when the test fails, it logs only a generic error message
(see [build
log](https://buildkite.com/ray-project/postmerge/builds/10962/steps/canvas?jid=019788f1-08fe-4509-a42e-85cd1f1b6e37#019788f1-08fe-4509-a42e-85cd1f1b6e37/1546-1688)),
which makes debugging difficult. I'm replacing wait_for_condition with
wait_for_assertion, a more informative alternative that logs the actual
value being tested, making it easier to identify what went wrong.

Test:
 -CI

Signed-off-by: Cuong Nguyen <can@anyscale.com>

Author: can-anyscale

python/ray/tests/test_task_metrics.py

@@ -12,6 +12,7 @@
     raw_metrics,
     run_string_as_driver,
     run_string_as_driver_nonblocking,
+    wait_for_assertion,
 )
 
 
@@ -152,8 +153,14 @@ def f():
         "RUNNING_IN_RAY_GET": 1.0,
         "PENDING_NODE_ASSIGNMENT": 8.0,
     }
-    wait_for_condition(
-        lambda: tasks_by_state(info) == expected, timeout=20, retry_interval_ms=2000
+
+    def check_task_state():
+        assert tasks_by_state(info) == expected
+
+    wait_for_assertion(
+        check_task_state,
+        timeout=20,
+        retry_interval_ms=2000,
     )
     assert tasks_by_name_and_state(info) == {
         ("wrapper", "RUNNING_IN_RAY_GET"): 1.0,
@@ -190,8 +197,14 @@ def f():
         "RUNNING_IN_RAY_WAIT": 1.0,
         "PENDING_NODE_ASSIGNMENT": 8.0,
     }
-    wait_for_condition(
-        lambda: tasks_by_state(info) == expected, timeout=20, retry_interval_ms=2000
+
+    def check_task_state():
+        assert tasks_by_state(info) == expected
+
+    wait_for_assertion(
+        check_task_state,
+        timeout=20,
+        retry_interval_ms=2000,
     )
     assert tasks_by_name_and_state(info) == {
         ("wrapper", "RUNNING_IN_RAY_WAIT"): 1.0,