chore: import upstream snapshot with attribution

python/ray/tests/test_metric_cardinality.py

@@ -0,0 +1,155 @@
+# isort: skip_file
+# ruff: noqa: E402
+import sys
+import requests
+import os
+
+import pytest
+
+import ray
+from ray._private.test_utils import wait_for_assertion
+from ray._common.test_utils import (
+    fetch_prometheus_metric_timeseries,
+    PrometheusTimeseries,
+)
+from ray._common.network_utils import build_address
+from ray._private.telemetry.metric_cardinality import (
+    WORKER_ID_TAG_KEY,
+    TASK_OR_ACTOR_NAME_TAG_KEY,
+    MetricCardinality,
+)
+
+
+try:
+    import prometheus_client
+except ImportError:
+    prometheus_client = None
+
+
+_TO_TEST_METRICS = ["tasks", "actors", "running_jobs"]
+_COMPONENT_TAG_KEY = "Component"
+
+
+@pytest.fixture
+def _setup_cluster_for_test(request, ray_start_cluster):
+    global _CARDINALITY_LEVEL
+    _CARDINALITY_LEVEL = None
+    core_metric_cardinality_level = request.param
+    os.environ["RAY_metric_cardinality_level"] = core_metric_cardinality_level
+    cluster = ray_start_cluster
+    cluster.add_node(
+        _system_config={
+            "metrics_report_interval_ms": 1000,
+            "enable_metrics_collection": True,
+            "metric_cardinality_level": core_metric_cardinality_level,
+        }
+    )
+    cluster.wait_for_nodes()
+    ray_context = ray.init(
+        address=cluster.address,
+    )
+
+    @ray.remote
+    def t():
+        print("task")
+
+    @ray.remote
+    class A:
+        async def run(self):
+            print("actor")
+
+    a = A.remote()
+    obj_refs = [t.remote(), a.run.remote()]
+
+    # Make a request to the dashboard to produce some dashboard metrics
+    requests.get(f"http://{ray_context.dashboard_url}/nodes")
+
+    node_info_list = ray.nodes()
+    prom_addresses = []
+    for node_info in node_info_list:
+        prom_addresses.append(
+            build_address(
+                node_info["NodeManagerAddress"], node_info["MetricsExportPort"]
+            )
+        )
+    yield prom_addresses
+
+    ray.get(obj_refs)
+
+
+def _cardinality_level_test(_setup_cluster_for_test, cardinality_level, metric):
+    """
+    Test that the ray_tasks and ray_actors metric are reported with the expected cardinality level
+    """
+    TEST_TIMEOUT_S = 30
+    prom_addresses = _setup_cluster_for_test
+
+    def _validate():
+        timeseries = PrometheusTimeseries()
+        metric_samples = fetch_prometheus_metric_timeseries(prom_addresses, timeseries)
+        samples = metric_samples.get(f"ray_{metric}")
+        assert samples, f"Metric {metric} not found in samples"
+        for sample in samples:
+            if (
+                cardinality_level == "legacy"
+                or metric not in MetricCardinality.get_high_cardinality_metrics()
+            ):
+                # If the cardinality level is legacy, the WorkerId tag should be
+                # present
+                assert (
+                    sample.labels.get(WORKER_ID_TAG_KEY) is not None
+                ), f"Sample {sample} does not contain WorkerId tag"
+                if metric == "tasks" or metric == "actors":
+                    assert (
+                        sample.labels.get(TASK_OR_ACTOR_NAME_TAG_KEY) is not None
+                    ), f"Sample {sample} does not contain Name tag"
+            elif cardinality_level == "recommended":
+                # If the cardinality level is recommended, the WorkerId tag should
+                # be removed
+                assert (
+                    sample.labels.get(WORKER_ID_TAG_KEY) is None
+                ), f"Sample {sample} contains WorkerId tag"
+            elif cardinality_level == "low":
+                # If the cardinality level is low, the WorkerId and Name tags should
+                # be removed
+                assert (
+                    sample.labels.get(WORKER_ID_TAG_KEY) is None
+                ), f"Sample {sample} contains WorkerId tag"
+                assert (
+                    sample.labels.get(TASK_OR_ACTOR_NAME_TAG_KEY) is None
+                ), f"Sample {sample} contains Name tag"
+            else:
+                raise ValueError(f"Unknown cardinality level: {cardinality_level}")
+
+            # The Component tag should be present on all cardinality levels
+            assert (
+                sample.labels.get(_COMPONENT_TAG_KEY) is not None
+            ), f"Sample {sample} does not contain Component tag"
+
+    wait_for_assertion(
+        _validate,
+        timeout=TEST_TIMEOUT_S,
+        retry_interval_ms=1000,  # Yield resource for other processes
+    )
+
+
+@pytest.mark.skipif(prometheus_client is None, reason="Prometheus not installed")
+@pytest.mark.parametrize(
+    "_setup_cluster_for_test,cardinality_level,metric",
+    [
+        (cardinality, cardinality, metric)
+        for cardinality in ["low", "recommended", "legacy"]
+        for metric in _TO_TEST_METRICS
+    ],
+    indirect=["_setup_cluster_for_test"],
+)
+def test_cardinality_recommended_and_legacy_levels(
+    _setup_cluster_for_test, cardinality_level, metric
+):
+    _cardinality_level_test(_setup_cluster_for_test, cardinality_level, metric)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-sv", __file__]))