[core] deprecate raw_metrics api in tests (#58614)

Deprecate `raw_metrics` API and replace them all with the
`raw_metric_timeseries` API. `raw_metrics` returns the current snapshot
of a set of metrics, while `raw_metric_timeseries` returns the full time
series. The later is more reliable when checking the latest instance of
several independent metrics.

Test:
- CI

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

Author: can-anyscale

python/ray/_private/test_utils.py

@@ -1065,20 +1065,6 @@ def fetch_prometheus_metric_timeseries(
     return samples_by_name
 
 
-def raw_metrics(info: RayContext) -> Dict[str, List[Any]]:
-    """Return prometheus metrics from a RayContext
-
-    Args:
-        info: Ray context returned from ray.init()
-
-    Returns:
-        Dict from metric name to a list of samples for the metrics
-    """
-    metrics_page = "localhost:{}".format(info.address_info["metrics_export_port"])
-    print("Fetch metrics from", metrics_page)
-    return fetch_prometheus_metrics([metrics_page])
-
-
 def raw_metric_timeseries(
     info: RayContext, result: PrometheusTimeseries
 ) -> Dict[str, List[Any]]:

python/ray/tests/test_memory_pressure.py

@@ -13,7 +13,7 @@
 )
 from ray._private.grpc_utils import init_grpc_channel
 from ray._private.state_api_test_utils import verify_failed_task
-from ray._private.test_utils import raw_metrics
+from ray._private.test_utils import PrometheusTimeseries, raw_metric_timeseries
 from ray._private.utils import get_used_memory
 from ray.util.state.state_manager import StateDataSourceClient
 
@@ -118,8 +118,10 @@ def get_additional_bytes_to_reach_memory_usage_pct(pct: float) -> int:
     return bytes_needed
 
 
-def has_metric_tagged_with_value(addr, tag, value) -> bool:
-    metrics = raw_metrics(addr)
+def has_metric_tagged_with_value(
+    addr, tag, value, timeseries: PrometheusTimeseries
+) -> bool:
+    metrics = raw_metric_timeseries(addr, timeseries)
     for name, samples in metrics.items():
         for sample in samples:
             if tag in set(sample.labels.values()) and sample.value == value:
@@ -145,13 +147,15 @@ def test_restartable_actor_throws_oom_error(ray_with_memory_monitor, restartable
     with pytest.raises(ray.exceptions.OutOfMemoryError):
         ray.get(leaker.allocate.remote(bytes_to_alloc, memory_monitor_refresh_ms * 3))
 
+    timeseries = PrometheusTimeseries()
     wait_for_condition(
         has_metric_tagged_with_value,
         timeout=10,
         retry_interval_ms=100,
         addr=addr,
         tag="MemoryManager.ActorEviction.Total",
         value=2.0 if restartable else 1.0,
+        timeseries=timeseries,
     )
 
     wait_for_condition(
@@ -161,6 +165,7 @@ def test_restartable_actor_throws_oom_error(ray_with_memory_monitor, restartable
         addr=addr,
         tag="Leaker.__init__",
         value=2.0 if restartable else 1.0,
+        timeseries=timeseries,
     )
 
 
@@ -180,13 +185,15 @@ def test_restartable_actor_oom_retry_off_throws_oom_error(
     with pytest.raises(ray.exceptions.OutOfMemoryError) as _:
         ray.get(leaker.allocate.remote(bytes_to_alloc, memory_monitor_refresh_ms * 3))
 
+    timeseries = PrometheusTimeseries()
     wait_for_condition(
         has_metric_tagged_with_value,
         timeout=10,
         retry_interval_ms=100,
         addr=addr,
         tag="MemoryManager.ActorEviction.Total",
         value=2.0,
+        timeseries=timeseries,
     )
     wait_for_condition(
         has_metric_tagged_with_value,
@@ -195,6 +202,7 @@ def test_restartable_actor_oom_retry_off_throws_oom_error(
         addr=addr,
         tag="Leaker.__init__",
         value=2.0,
+        timeseries=timeseries,
     )
 
 
@@ -210,13 +218,15 @@ def test_non_retryable_task_killed_by_memory_monitor_with_oom_error(
     with pytest.raises(ray.exceptions.OutOfMemoryError) as _:
         ray.get(allocate_memory.options(max_retries=0).remote(bytes_to_alloc))
 
+    timeseries = PrometheusTimeseries()
     wait_for_condition(
         has_metric_tagged_with_value,
         timeout=10,
         retry_interval_ms=100,
         addr=addr,
         tag="MemoryManager.TaskEviction.Total",
         value=1.0,
+        timeseries=timeseries,
     )
     wait_for_condition(
         has_metric_tagged_with_value,
@@ -225,6 +235,7 @@ def test_non_retryable_task_killed_by_memory_monitor_with_oom_error(
         addr=addr,
         tag="allocate_memory",
         value=1.0,
+        timeseries=timeseries,
     )
 
 
@@ -372,13 +383,15 @@ def test_task_oom_no_oom_retry_fails_immediately(
             )
         )
 
+    timeseries = PrometheusTimeseries()
     wait_for_condition(
         has_metric_tagged_with_value,
         timeout=10,
         retry_interval_ms=100,
         addr=addr,
         tag="MemoryManager.TaskEviction.Total",
         value=1.0,
+        timeseries=timeseries,
     )
     wait_for_condition(
         has_metric_tagged_with_value,
@@ -387,6 +400,7 @@ def test_task_oom_no_oom_retry_fails_immediately(
         addr=addr,
         tag="allocate_memory",
         value=1.0,
+        timeseries=timeseries,
     )
 
 
@@ -411,13 +425,15 @@ def test_task_oom_only_uses_oom_retry(
             )
         )
 
+    timeseries = PrometheusTimeseries()
     wait_for_condition(
         has_metric_tagged_with_value,
         timeout=10,
         retry_interval_ms=100,
         addr=addr,
         tag="MemoryManager.TaskEviction.Total",
         value=task_oom_retries + 1,
+        timeseries=timeseries,
     )
     wait_for_condition(
         has_metric_tagged_with_value,
@@ -426,6 +442,7 @@ def test_task_oom_only_uses_oom_retry(
         addr=addr,
         tag="allocate_memory",
         value=task_oom_retries + 1,
+        timeseries=timeseries,
     )
 
 
@@ -502,6 +519,7 @@ def infinite_retry_task():
             time.sleep(5)
 
     with ray.init() as addr:
+        timeseries = PrometheusTimeseries()
         with pytest.raises(ray.exceptions.OutOfMemoryError) as _:
             ray.get(infinite_retry_task.remote())
 
@@ -512,6 +530,7 @@ def infinite_retry_task():
             addr=addr,
             tag="MemoryManager.TaskEviction.Total",
             value=1.0,
+            timeseries=timeseries,
         )
 
 

python/ray/tests/test_object_store_metrics.py

@@ -9,7 +9,8 @@
 import ray
 from ray._common.test_utils import wait_for_condition
 from ray._private.test_utils import (
-    raw_metrics,
+    PrometheusTimeseries,
+    raw_metric_timeseries,
 )
 from ray._private.worker import RayContext
 from ray.dashboard.consts import RAY_DASHBOARD_STATS_UPDATING_INTERVAL
@@ -26,8 +27,10 @@
 }
 
 
-def _objects_by_tag(info: RayContext, tag: str) -> Dict:
-    res = raw_metrics(info)
+def _objects_by_tag(
+    info: RayContext, tag: str, timeseries: PrometheusTimeseries
+) -> Dict:
+    res = raw_metric_timeseries(info, timeseries)
     objects_info = defaultdict(int)
     if "ray_object_store_memory" in res:
         for sample in res["ray_object_store_memory"]:
@@ -41,12 +44,12 @@ def _objects_by_tag(info: RayContext, tag: str) -> Dict:
     return objects_info
 
 
-def objects_by_seal_state(info: RayContext) -> Dict:
-    return _objects_by_tag(info, "ObjectState")
+def objects_by_seal_state(info: RayContext, timeseries: PrometheusTimeseries) -> Dict:
+    return _objects_by_tag(info, "ObjectState", timeseries)
 
 
-def objects_by_loc(info: RayContext) -> Dict:
-    return _objects_by_tag(info, "Location")
+def objects_by_loc(info: RayContext, timeseries: PrometheusTimeseries) -> Dict:
+    return _objects_by_tag(info, "Location", timeseries)
 
 
 def approx_eq_dict_in(actual: Dict, expected: Dict, e: int) -> bool:
@@ -79,6 +82,7 @@ def test_shared_memory_and_inline_worker_heap(shutdown_only):
             },
         },
     )
+    timeseries = PrometheusTimeseries()
 
     # Allocate 80MiB data
     objs_in_use = ray.get(
@@ -94,7 +98,7 @@ def test_shared_memory_and_inline_worker_heap(shutdown_only):
 
     wait_for_condition(
         # 1KiB for metadata difference
-        lambda: approx_eq_dict_in(objects_by_loc(info), expected, 2 * KiB),
+        lambda: approx_eq_dict_in(objects_by_loc(info, timeseries), expected, 2 * KiB),
         timeout=20,
         retry_interval_ms=500,
     )
@@ -117,7 +121,7 @@ def func():
 
     wait_for_condition(
         # 4 KiB for metadata difference
-        lambda: approx_eq_dict_in(objects_by_loc(info), expected, 4 * KiB),
+        lambda: approx_eq_dict_in(objects_by_loc(info, timeseries), expected, 4 * KiB),
         timeout=20,
         retry_interval_ms=500,
     )
@@ -136,7 +140,7 @@ def func():
 
     wait_for_condition(
         # 1KiB for metadata difference
-        lambda: approx_eq_dict_in(objects_by_loc(info), expected, 2 * KiB),
+        lambda: approx_eq_dict_in(objects_by_loc(info, timeseries), expected, 2 * KiB),
         timeout=20,
         retry_interval_ms=500,
     )
@@ -158,7 +162,7 @@ def test_spilling(object_spilling_config, shutdown_only):
             **{"object_spilling_config": object_spilling_config},
         },
     )
-
+    timeseries = PrometheusTimeseries()
     # Create and use 100MiB data, which should fit in memory
     objs1 = [ray.put(np.zeros(50 * MiB, dtype=np.uint8)) for _ in range(2)]
 
@@ -171,7 +175,7 @@ def test_spilling(object_spilling_config, shutdown_only):
 
     wait_for_condition(
         # 1KiB for metadata difference
-        lambda: approx_eq_dict_in(objects_by_loc(info), expected, 1 * KiB),
+        lambda: approx_eq_dict_in(objects_by_loc(info, timeseries), expected, 1 * KiB),
         timeout=20,
         retry_interval_ms=500,
     )
@@ -187,7 +191,7 @@ def test_spilling(object_spilling_config, shutdown_only):
     }
     wait_for_condition(
         # 1KiB for metadata difference
-        lambda: approx_eq_dict_in(objects_by_loc(info), expected, 1 * KiB),
+        lambda: approx_eq_dict_in(objects_by_loc(info, timeseries), expected, 1 * KiB),
         timeout=20,
         retry_interval_ms=500,
     )
@@ -202,7 +206,7 @@ def test_spilling(object_spilling_config, shutdown_only):
     }
     wait_for_condition(
         # 1KiB for metadata difference
-        lambda: approx_eq_dict_in(objects_by_loc(info), expected, 1 * KiB),
+        lambda: approx_eq_dict_in(objects_by_loc(info, timeseries), expected, 1 * KiB),
         timeout=20,
         retry_interval_ms=500,
     )
@@ -217,7 +221,7 @@ def test_spilling(object_spilling_config, shutdown_only):
     }
     wait_for_condition(
         # 1KiB for metadata difference
-        lambda: approx_eq_dict_in(objects_by_loc(info), expected, 1 * KiB),
+        lambda: approx_eq_dict_in(objects_by_loc(info, timeseries), expected, 1 * KiB),
         timeout=20,
         retry_interval_ms=500,
     )
@@ -239,6 +243,7 @@ def test_fallback_memory(shutdown_only):
         object_store_memory=expected_in_memory * obj_size_mb * MiB + delta_mb * MiB,
         _system_config=_SYSTEM_CONFIG,
     )
+    timeseries = PrometheusTimeseries()
     obj_refs = [
         ray.put(np.zeros(obj_size_mb * MiB, dtype=np.uint8))
         for _ in range(expected_in_memory)
@@ -257,7 +262,7 @@ def test_fallback_memory(shutdown_only):
 
     wait_for_condition(
         # 2KiB for metadata difference
-        lambda: approx_eq_dict_in(objects_by_loc(info), expected, 3 * KiB),
+        lambda: approx_eq_dict_in(objects_by_loc(info, timeseries), expected, 3 * KiB),
         timeout=20,
         retry_interval_ms=500,
     )
@@ -285,7 +290,7 @@ def test_fallback_memory(shutdown_only):
 
     wait_for_condition(
         # 3KiB for metadata difference
-        lambda: approx_eq_dict_in(objects_by_loc(info), expected, 3 * KiB),
+        lambda: approx_eq_dict_in(objects_by_loc(info, timeseries), expected, 3 * KiB),
         timeout=20,
         retry_interval_ms=500,
     )
@@ -305,7 +310,7 @@ def test_fallback_memory(shutdown_only):
 
     wait_for_condition(
         # 3KiB for metadata difference
-        lambda: approx_eq_dict_in(objects_by_loc(info), expected, 3 * KiB),
+        lambda: approx_eq_dict_in(objects_by_loc(info, timeseries), expected, 3 * KiB),
         timeout=20,
         retry_interval_ms=500,
     )
@@ -322,7 +327,7 @@ def test_seal_memory(shutdown_only):
         object_store_memory=100 * MiB,
         _system_config=_SYSTEM_CONFIG,
     )
-
+    timeseries = PrometheusTimeseries()
     # Allocate 80MiB data
     objs_in_use = ray.get(
         [ray.put(np.zeros(20 * MiB, dtype=np.uint8)) for _ in range(4)]
@@ -335,7 +340,9 @@ def test_seal_memory(shutdown_only):
 
     wait_for_condition(
         # 1KiB for metadata difference
-        lambda: approx_eq_dict_in(objects_by_seal_state(info), expected, 2 * KiB),
+        lambda: approx_eq_dict_in(
+            objects_by_seal_state(info, timeseries), expected, 2 * KiB
+        ),
         timeout=20,
         retry_interval_ms=500,
     )
@@ -349,7 +356,9 @@ def test_seal_memory(shutdown_only):
 
     wait_for_condition(
         # 1KiB for metadata difference
-        lambda: approx_eq_dict_in(objects_by_seal_state(info), expected, 2 * KiB),
+        lambda: approx_eq_dict_in(
+            objects_by_seal_state(info, timeseries), expected, 2 * KiB
+        ),
         timeout=20,
         retry_interval_ms=500,
     )
@@ -362,9 +371,10 @@ def test_object_store_memory_matches_dashboard_obj_memory(shutdown_only):
     ctx = ray.init(
         object_store_memory=500 * MiB,
     )
+    timeseries = PrometheusTimeseries()
 
     def verify():
-        resources = raw_metrics(ctx)["ray_resources"]
+        resources = raw_metric_timeseries(ctx, timeseries)["ray_resources"]
         object_store_memory_bytes_from_metrics = 0
         for sample in resources:
             # print(sample)