chore: import upstream snapshot with attribution

python/ray/dashboard/modules/data/data_head.py

@@ -0,0 +1,161 @@
+import json
+import logging
+import os
+from enum import Enum
+from urllib.parse import quote
+
+import aiohttp
+from aiohttp.web import Request, Response
+
+import ray.dashboard.optional_utils as optional_utils
+from ray.dashboard.modules.metrics.metrics_head import (
+    DEFAULT_PROMETHEUS_HEADERS,
+    DEFAULT_PROMETHEUS_HOST,
+    PROMETHEUS_HEADERS_ENV_VAR,
+    PROMETHEUS_HOST_ENV_VAR,
+    PrometheusQueryError,
+    parse_prom_headers,
+)
+from ray.dashboard.subprocesses.module import SubprocessModule
+from ray.dashboard.subprocesses.routes import SubprocessRouteTable as routes
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+
+# Window and sampling rate used for certain Prometheus queries.
+# Datapoints up until `MAX_TIME_WINDOW` ago are queried at `SAMPLE_RATE` intervals.
+MAX_TIME_WINDOW = "1h"
+SAMPLE_RATE = "1s"
+
+
+class PrometheusQuery(Enum):
+    """Enum to store types of Prometheus queries for a given metric and grouping."""
+
+    VALUE = ("value", "sum({}{{SessionName='{}'}}) by ({})")
+    MAX = (
+        "max",
+        "max_over_time(sum({}{{SessionName='{}'}}) by ({})["
+        + f"{MAX_TIME_WINDOW}:{SAMPLE_RATE}])",
+    )
+
+
+DATASET_METRICS = {
+    "ray_data_output_rows": (PrometheusQuery.MAX,),
+    "ray_data_spilled_bytes": (PrometheusQuery.MAX,),
+    "ray_data_current_bytes": (PrometheusQuery.VALUE, PrometheusQuery.MAX),
+    "ray_data_cpu_usage_cores": (PrometheusQuery.VALUE, PrometheusQuery.MAX),
+    "ray_data_gpu_usage_cores": (PrometheusQuery.VALUE, PrometheusQuery.MAX),
+}
+
+
+class DataHead(SubprocessModule):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.prometheus_host = os.environ.get(
+            PROMETHEUS_HOST_ENV_VAR, DEFAULT_PROMETHEUS_HOST
+        )
+        self.prometheus_headers = parse_prom_headers(
+            os.environ.get(
+                PROMETHEUS_HEADERS_ENV_VAR,
+                DEFAULT_PROMETHEUS_HEADERS,
+            )
+        )
+
+    @routes.get("/api/data/datasets/{job_id}")
+    @optional_utils.init_ray_and_catch_exceptions()
+    async def get_datasets(self, req: Request) -> Response:
+        job_id = req.match_info["job_id"]
+
+        try:
+            from ray.data._internal.stats import get_or_create_stats_actor
+
+            _stats_actor = get_or_create_stats_actor()
+            datasets = await _stats_actor.get_datasets.remote(job_id)
+            # Initializes dataset metric values
+            for dataset in datasets:
+                for metric, queries in DATASET_METRICS.items():
+                    datasets[dataset][metric] = {query.value[0]: 0 for query in queries}
+                    for operator in datasets[dataset]["operators"]:
+                        datasets[dataset]["operators"][operator][metric] = {
+                            query.value[0]: 0 for query in queries
+                        }
+            # Query dataset metric values from prometheus
+            try:
+                # TODO (Zandew): store results of completed datasets in stats actor.
+                for metric, queries in DATASET_METRICS.items():
+                    for query in queries:
+                        query_name, prom_query = query.value
+                        # Dataset level
+                        dataset_result = await self._query_prometheus(
+                            prom_query.format(metric, self.session_name, "dataset")
+                        )
+                        for res in dataset_result["data"]["result"]:
+                            dataset, value = res["metric"]["dataset"], res["value"][1]
+                            if dataset in datasets:
+                                datasets[dataset][metric][query_name] = value
+
+                        # Operator level
+                        operator_result = await self._query_prometheus(
+                            prom_query.format(
+                                metric, self.session_name, "dataset, operator"
+                            )
+                        )
+                        for res in operator_result["data"]["result"]:
+                            dataset, operator, value = (
+                                res["metric"]["dataset"],
+                                res["metric"]["operator"],
+                                res["value"][1],
+                            )
+                            # Check if dataset/operator is in current _StatsActor scope.
+                            # Prometheus server may contain metrics from previous
+                            # cluster if not reset.
+                            if (
+                                dataset in datasets
+                                and operator in datasets[dataset]["operators"]
+                            ):
+                                datasets[dataset]["operators"][operator][metric][
+                                    query_name
+                                ] = value
+            except aiohttp.client_exceptions.ClientConnectorError:
+                # Prometheus server may not be running,
+                # leave these values blank and return other data
+                logging.exception(
+                    "Exception occurred while querying Prometheus. "
+                    "The Prometheus server may not be running."
+                )
+            # Flatten response
+            for dataset in datasets:
+                datasets[dataset]["operators"] = list(
+                    map(
+                        lambda item: {"operator": item[0], **item[1]},
+                        datasets[dataset]["operators"].items(),
+                    )
+                )
+            datasets = list(
+                map(lambda item: {"dataset": item[0], **item[1]}, datasets.items())
+            )
+            # Sort by descending start time
+            datasets = sorted(datasets, key=lambda x: x["start_time"], reverse=True)
+            return Response(
+                text=json.dumps({"datasets": datasets}),
+                content_type="application/json",
+            )
+        except Exception as e:
+            logging.exception("Exception occurred while getting datasets.")
+            return Response(
+                status=503,
+                text=str(e),
+            )
+
+    async def _query_prometheus(self, query):
+        async with self.http_session.get(
+            f"{self.prometheus_host}/api/v1/query?query={quote(query)}",
+            headers=self.prometheus_headers,
+        ) as resp:
+            if resp.status == 200:
+                prom_data = await resp.json()
+                return prom_data
+
+            message = await resp.text()
+            raise PrometheusQueryError(resp.status, message)

python/ray/dashboard/modules/data/tests/test_data_head.py

@@ -0,0 +1,141 @@
+import os
+import sys
+
+import pytest
+import requests
+
+import ray
+from ray.job_submission import JobSubmissionClient
+from ray.tests.conftest import *  # noqa
+
+# For local testing on a Macbook, set `export TEST_ON_DARWIN=1`.
+TEST_ON_DARWIN = os.environ.get("TEST_ON_DARWIN", "0") == "1"
+
+DATA_HEAD_URLS = {"GET": "http://localhost:8265/api/data/datasets/{job_id}"}
+
+DATA_SCHEMA = [
+    "state",
+    "progress",
+    "total",
+    "total_rows",
+    "ray_data_output_rows",
+    "ray_data_spilled_bytes",
+    "ray_data_current_bytes",
+    "ray_data_cpu_usage_cores",
+    "ray_data_gpu_usage_cores",
+]
+
+RESPONSE_SCHEMA = [
+    "dataset",
+    "job_id",
+    "start_time",
+    "end_time",
+    "operators",
+] + DATA_SCHEMA
+
+OPERATOR_SCHEMA = [
+    "name",
+    "operator",
+    "queued_blocks",
+] + DATA_SCHEMA
+
+
+@pytest.mark.skipif(
+    sys.platform == "darwin" and not TEST_ON_DARWIN, reason="Flaky on OSX."
+)
+def test_unique_operator_id(ray_start_regular_shared):
+    # This regression test addresses a bug caused by using a non-unique operator ID
+    # format. Specifically, the third operator's name is limit11 with the ID limit112,
+    # while the thirteenth operator's name is limit1 with the same ID limit112, leading
+    # to a collision.
+    ds = ray.data.range(100, override_num_blocks=20).limit(11)  # 3 operators
+    for i in range(11):  # 11 more operators
+        ds = ds.limit(1)
+    ds._set_name("unique_operator_id_test")
+    ds.materialize()
+
+    client = JobSubmissionClient()
+    jobs = client.list_jobs()
+    assert len(jobs) == 1, jobs
+    job_id = jobs[0].job_id
+
+    data = requests.get(DATA_HEAD_URLS["GET"].format(job_id=job_id)).json()
+    datasets = [
+        dataset
+        for dataset in data["datasets"]
+        if dataset["dataset"].startswith("unique_operator_id_test")
+    ]
+    assert len(datasets) == 1
+    dataset = datasets[0]
+
+    operators = dataset["operators"]
+    assert len(operators) == 3  # Should be 3 because of limiter operator fusion.
+
+
+@pytest.mark.skipif(
+    sys.platform == "darwin" and not TEST_ON_DARWIN, reason="Flaky on OSX."
+)
+def test_get_datasets(ray_start_regular_shared):
+    ds = ray.data.range(100, override_num_blocks=20).map_batches(lambda x: x)
+    ds.set_name("data_head_test")
+    ds.materialize()
+
+    client = JobSubmissionClient()
+    jobs = client.list_jobs()
+    assert len(jobs) == 1, jobs
+    job_id = jobs[0].job_id
+
+    data = requests.get(DATA_HEAD_URLS["GET"].format(job_id=job_id)).json()
+    datasets = [
+        dataset
+        for dataset in data["datasets"]
+        if dataset["dataset"].startswith("data_head_test")
+    ]
+
+    assert len(datasets) == 1
+    assert sorted(datasets[0].keys()) == sorted(RESPONSE_SCHEMA)
+
+    dataset = datasets[0]
+    assert dataset["dataset"].startswith("data_head_test")
+    assert dataset["job_id"] == job_id
+    assert dataset["state"] == "FINISHED"
+    assert dataset["end_time"] is not None
+
+    operators = dataset["operators"]
+    assert len(operators) == 2
+    op0 = operators[0]
+    op1 = operators[1]
+    assert sorted(op0.keys()) == sorted(OPERATOR_SCHEMA)
+    assert sorted(op1.keys()) == sorted(OPERATOR_SCHEMA)
+    assert {
+        "operator": "Input_0",
+        "name": "Input",
+        "state": "FINISHED",
+        "progress": 20,
+        "total": 20,
+    }.items() <= op0.items()
+    assert {
+        "operator": "ReadRange->MapBatches(<lambda>)_1",
+        "name": "ReadRange->MapBatches(<lambda>)",
+        "state": "FINISHED",
+        "progress": 20,
+        "total": 20,
+    }.items() <= op1.items()
+
+    ds._set_name("another_data_head_test")
+    ds.map_batches(lambda x: x).materialize()
+    data = requests.get(DATA_HEAD_URLS["GET"].format(job_id=job_id)).json()
+
+    dataset = [
+        dataset
+        for dataset in data["datasets"]
+        if dataset["dataset"].startswith("another_data_head_test")
+    ][0]
+    assert dataset["dataset"].startswith("another_data_head_test")
+    assert dataset["job_id"] == job_id
+    assert dataset["state"] == "FINISHED"
+    assert dataset["end_time"] is not None
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-vv", __file__]))