wehub-resource-sync
142 lines
4.0 KiB
Python
142 lines
4.0 KiB
Python
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__]))
|