chore: import upstream snapshot with attribution

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__]))