wehub-resource-sync
762 lines
26 KiB
Python
762 lines
26 KiB
Python
import json
|
|
import os
|
|
from dataclasses import asdict, dataclass, field
|
|
from typing import Any, Dict, List, Tuple
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
from ray.data._internal.execution.dataset_state import DatasetState
|
|
from ray.data._internal.logical.interfaces import LogicalOperator
|
|
from ray.data._internal.metadata_exporter import (
|
|
UNKNOWN,
|
|
DataContextMetadata,
|
|
Operator,
|
|
Topology,
|
|
sanitize_for_struct,
|
|
)
|
|
from ray.data._internal.stats import get_or_create_stats_actor
|
|
from ray.data.context import DataContext
|
|
from ray.tests.conftest import _ray_start
|
|
|
|
STUB_JOB_ID = "stub_job_id"
|
|
STUB_DATASET_ID = "stub_dataset_id"
|
|
|
|
|
|
def _get_export_file_path() -> str:
|
|
return os.path.join(
|
|
ray._private.worker._global_node.get_session_dir_path(),
|
|
"logs",
|
|
"export_events",
|
|
"event_EXPORT_DATASET_METADATA.log",
|
|
)
|
|
|
|
|
|
def _get_exported_data():
|
|
exported_file = _get_export_file_path()
|
|
assert os.path.isfile(exported_file)
|
|
|
|
with open(exported_file, "r") as f:
|
|
data = f.readlines()
|
|
|
|
return [json.loads(line) for line in data]
|
|
|
|
|
|
@pytest.fixture
|
|
def ray_start_cluster_with_export_api_config(shutdown_only):
|
|
"""Enable export API for the EXPORT_TRAIN_RUN resource type."""
|
|
with _ray_start(
|
|
num_cpus=4,
|
|
runtime_env={
|
|
"env_vars": {
|
|
"RAY_enable_export_api_write_config": "EXPORT_DATASET_METADATA"
|
|
}
|
|
},
|
|
) as res:
|
|
yield res
|
|
|
|
|
|
@pytest.fixture
|
|
def ray_start_cluster_with_export_api_write(shutdown_only):
|
|
"""Enable export API for all resource types."""
|
|
with _ray_start(
|
|
num_cpus=4,
|
|
runtime_env={"env_vars": {"RAY_enable_export_api_write": "1"}},
|
|
) as res:
|
|
yield res
|
|
|
|
|
|
@dataclass
|
|
class TestDataclass:
|
|
"""A test dataclass for testing dataclass serialization."""
|
|
|
|
list_field: list = None
|
|
dict_field: dict = None
|
|
string_field: str = "test"
|
|
int_field: int = 1
|
|
float_field: float = 1.0
|
|
set_field: set = None
|
|
tuple_field: Tuple[int] = None
|
|
bool_field: bool = True
|
|
none_field: None = None
|
|
|
|
def __post_init__(self):
|
|
self.list_field = [1, 2, 3]
|
|
self.dict_field = {1: 2, "3": "4"}
|
|
self.set_field = {1, 2, 3}
|
|
self.tuple_field = (1, 2, 3)
|
|
|
|
|
|
@dataclass(frozen=True, repr=False, eq=False)
|
|
class DummyLogicalOperator(LogicalOperator):
|
|
"""A dummy logical operator for testing _get_logical_args with various data types."""
|
|
|
|
_name: str = field(init=False, default="DummyOperator", repr=False)
|
|
_input_dependencies: List[LogicalOperator] = field(
|
|
init=False, default_factory=list, repr=False
|
|
)
|
|
_num_outputs: None = field(init=False, default=None, repr=False)
|
|
_string_value: str = "test_string"
|
|
_int_value: int = 42
|
|
_float_value: float = 3.14
|
|
_bool_value: bool = True
|
|
_none_value: None = None
|
|
_list_value: List[Any] = field(default_factory=lambda: [1, 2, 3, "string", None])
|
|
_dict_value: Dict[str, Any] = field(
|
|
default_factory=lambda: {"key1": "value1", "key2": 123, "key3": None}
|
|
)
|
|
_nested_dict: Dict[str, Any] = field(
|
|
default_factory=lambda: {
|
|
"level1": {
|
|
"level2": {
|
|
"level3": "deep_value",
|
|
"numbers": [1, 2, 3],
|
|
"mixed": {"a": 1, "b": "string", "c": None},
|
|
}
|
|
}
|
|
}
|
|
)
|
|
_tuple_value: Tuple[Any, ...] = (1, "string", None, 3.14)
|
|
_set_value: set = field(default_factory=lambda: {1})
|
|
_bytes_value: bytes = b"binary_data"
|
|
_complex_dict: Dict[str, Any] = field(
|
|
default_factory=lambda: {
|
|
"string_keys": {"a": 1, "b": 2},
|
|
"int_keys": {
|
|
1: "one",
|
|
2: "two",
|
|
}, # This should cause issues if not handled
|
|
"mixed_keys": {"str": "value", 1: "int_key", None: "none_key"},
|
|
}
|
|
)
|
|
_empty_containers: Dict[str, Any] = field(
|
|
default_factory=lambda: {
|
|
"empty_list": [],
|
|
"empty_dict": {},
|
|
"empty_tuple": (),
|
|
"empty_set": set(),
|
|
}
|
|
)
|
|
_special_values: Dict[str, Any] = field(
|
|
default_factory=lambda: {
|
|
"zero": 0,
|
|
"negative": -1,
|
|
"large_int": 999999999999999999,
|
|
"small_float": 0.0000001,
|
|
"inf": float("inf"),
|
|
"neg_inf": float("-inf"),
|
|
"nan": float("nan"),
|
|
}
|
|
)
|
|
_data_class: TestDataclass = field(default_factory=TestDataclass)
|
|
|
|
@property
|
|
def num_outputs(self):
|
|
return self._num_outputs
|
|
|
|
|
|
@pytest.fixture
|
|
def dummy_dataset_topology():
|
|
"""Create a dummy Topology."""
|
|
dummy_operator = DummyLogicalOperator()
|
|
dummy_topology = Topology(
|
|
operators=[
|
|
Operator(
|
|
name="Input",
|
|
id="Input_0",
|
|
uuid="uuid_0",
|
|
input_dependencies=[],
|
|
sub_stages=[],
|
|
execution_start_time=1.0,
|
|
execution_end_time=1.0,
|
|
state="FINISHED",
|
|
args=sanitize_for_struct(dummy_operator._get_args()),
|
|
),
|
|
Operator(
|
|
name="ReadRange->Map(<lambda>)->Filter(<lambda>)",
|
|
id="ReadRange->Map(<lambda>)->Filter(<lambda>)_1",
|
|
uuid="uuid_1",
|
|
input_dependencies=["Input_0"],
|
|
sub_stages=[],
|
|
execution_start_time=0.0,
|
|
execution_end_time=0.0,
|
|
state="RUNNING",
|
|
args=sanitize_for_struct(dummy_operator._get_args()),
|
|
),
|
|
],
|
|
)
|
|
return dummy_topology
|
|
|
|
|
|
@pytest.fixture
|
|
def dummy_dataset_topology_expected_output():
|
|
return {
|
|
"operators": [
|
|
{
|
|
"name": "Input",
|
|
"id": "Input_0",
|
|
"uuid": "uuid_0",
|
|
"args": {
|
|
"_num_outputs": "None",
|
|
"_int_value": "42",
|
|
"_special_values": {
|
|
"negative": "-1",
|
|
"inf": "inf",
|
|
"zero": "0",
|
|
"large_int": "999999999999999999",
|
|
"small_float": "1e-07",
|
|
"neg_inf": "-inf",
|
|
"nan": "nan",
|
|
},
|
|
"_none_value": "None",
|
|
"_name": "DummyOperator",
|
|
"_output_dependencies": [],
|
|
"_float_value": "3.14",
|
|
"_list_value": ["1", "2", "3", "string", "None"],
|
|
"_dict_value": {"key1": "value1", "key3": "None", "key2": "123"},
|
|
"_set_value": ["1"],
|
|
"_tuple_value": ["1", "string", "None", "3.14"],
|
|
"_bytes_value": [
|
|
"98",
|
|
"105",
|
|
"110",
|
|
"97",
|
|
"114",
|
|
"121",
|
|
"95",
|
|
"100",
|
|
"97",
|
|
"116",
|
|
"97",
|
|
],
|
|
"_input_dependencies": [],
|
|
"_empty_containers": {
|
|
"empty_set": [],
|
|
"empty_tuple": [],
|
|
"empty_dict": {},
|
|
"empty_list": [],
|
|
},
|
|
"_bool_value": "True",
|
|
"_nested_dict": {
|
|
"level1": {
|
|
"level2": {
|
|
"mixed": {"a": "1", "b": "string", "c": "None"},
|
|
"numbers": ["1", "2", "3"],
|
|
"level3": "deep_value",
|
|
}
|
|
}
|
|
},
|
|
"_string_value": "test_string",
|
|
"_complex_dict": {
|
|
"string_keys": {"a": "1", "b": "2"},
|
|
"mixed_keys": {
|
|
"None": "none_key",
|
|
"str": "value",
|
|
"1": "int_key",
|
|
},
|
|
"int_keys": {"1": "one", "2": "two"},
|
|
},
|
|
"_data_class": {
|
|
"list_field": ["1", "2", "3"],
|
|
"dict_field": {"3": "4", "1": "2"},
|
|
"tuple_field": ["1", "2", "3"],
|
|
"set_field": ["1", "2", "3"],
|
|
"int_field": "1",
|
|
"none_field": "None",
|
|
"bool_field": "True",
|
|
"string_field": "test",
|
|
"float_field": "1.0",
|
|
},
|
|
},
|
|
"input_dependencies": [],
|
|
"sub_stages": [],
|
|
"execution_start_time": 1.0,
|
|
"execution_end_time": 1.0,
|
|
"state": "FINISHED",
|
|
},
|
|
{
|
|
"name": "ReadRange->Map(<lambda>)->Filter(<lambda>)",
|
|
"id": "ReadRange->Map(<lambda>)->Filter(<lambda>)_1",
|
|
"uuid": "uuid_1",
|
|
"input_dependencies": ["Input_0"],
|
|
"args": {
|
|
"_num_outputs": "None",
|
|
"_int_value": "42",
|
|
"_special_values": {
|
|
"negative": "-1",
|
|
"inf": "inf",
|
|
"zero": "0",
|
|
"large_int": "999999999999999999",
|
|
"small_float": "1e-07",
|
|
"neg_inf": "-inf",
|
|
"nan": "nan",
|
|
},
|
|
"_none_value": "None",
|
|
"_name": "DummyOperator",
|
|
"_output_dependencies": [],
|
|
"_float_value": "3.14",
|
|
"_list_value": ["1", "2", "3", "string", "None"],
|
|
"_dict_value": {"key1": "value1", "key3": "None", "key2": "123"},
|
|
"_set_value": ["1"],
|
|
"_tuple_value": ["1", "string", "None", "3.14"],
|
|
"_bytes_value": [
|
|
"98",
|
|
"105",
|
|
"110",
|
|
"97",
|
|
"114",
|
|
"121",
|
|
"95",
|
|
"100",
|
|
"97",
|
|
"116",
|
|
"97",
|
|
],
|
|
"_input_dependencies": [],
|
|
"_empty_containers": {
|
|
"empty_set": [],
|
|
"empty_tuple": [],
|
|
"empty_dict": {},
|
|
"empty_list": [],
|
|
},
|
|
"_bool_value": "True",
|
|
"_nested_dict": {
|
|
"level1": {
|
|
"level2": {
|
|
"mixed": {"a": "1", "b": "string", "c": "None"},
|
|
"numbers": ["1", "2", "3"],
|
|
"level3": "deep_value",
|
|
}
|
|
}
|
|
},
|
|
"_string_value": "test_string",
|
|
"_complex_dict": {
|
|
"string_keys": {"a": "1", "b": "2"},
|
|
"mixed_keys": {
|
|
"None": "none_key",
|
|
"str": "value",
|
|
"1": "int_key",
|
|
},
|
|
"int_keys": {"1": "one", "2": "two"},
|
|
},
|
|
"_data_class": {
|
|
"list_field": ["1", "2", "3"],
|
|
"dict_field": {"3": "4", "1": "2"},
|
|
"tuple_field": ["1", "2", "3"],
|
|
"set_field": ["1", "2", "3"],
|
|
"int_field": "1",
|
|
"none_field": "None",
|
|
"bool_field": "True",
|
|
"string_field": "test",
|
|
"float_field": "1.0",
|
|
},
|
|
},
|
|
"sub_stages": [],
|
|
"execution_start_time": 0.0,
|
|
"execution_end_time": 0.0,
|
|
"state": "RUNNING",
|
|
},
|
|
]
|
|
}
|
|
|
|
|
|
def test_export_disabled(ray_start_regular, dummy_dataset_topology):
|
|
"""Test that no export files are created when export API is disabled."""
|
|
stats_actor = get_or_create_stats_actor()
|
|
|
|
# Create or update train run
|
|
ray.get(
|
|
stats_actor.register_dataset.remote(
|
|
dataset_tag="test_dataset",
|
|
operator_tags=["ReadRange->Map(<lambda>)->Filter(<lambda>)"],
|
|
topology=dummy_dataset_topology,
|
|
job_id=STUB_JOB_ID,
|
|
data_context=DataContextMetadata.from_data_context(
|
|
DataContext.get_current()
|
|
),
|
|
)
|
|
)
|
|
|
|
# Check that no export files were created
|
|
assert not os.path.exists(_get_export_file_path())
|
|
|
|
|
|
def _test_dataset_metadata_export(topology, dummy_dataset_topology_expected_output):
|
|
"""Test that dataset metadata export events are written when export API is enabled."""
|
|
stats_actor = get_or_create_stats_actor()
|
|
|
|
# Simulate a dataset registration
|
|
ray.get(
|
|
stats_actor.register_dataset.remote(
|
|
dataset_tag=STUB_DATASET_ID,
|
|
operator_tags=["ReadRange->Map(<lambda>)->Filter(<lambda>)"],
|
|
topology=topology,
|
|
job_id=STUB_JOB_ID,
|
|
data_context=DataContextMetadata.from_data_context(
|
|
DataContext.get_current()
|
|
),
|
|
)
|
|
)
|
|
|
|
# Check that export files were created
|
|
data = _get_exported_data()
|
|
assert len(data) == 1
|
|
assert data[0]["source_type"] == "EXPORT_DATASET_METADATA"
|
|
assert data[0]["event_data"]["topology"] == dummy_dataset_topology_expected_output
|
|
assert data[0]["event_data"]["dataset_id"] == STUB_DATASET_ID
|
|
assert data[0]["event_data"]["job_id"] == STUB_JOB_ID
|
|
assert data[0]["event_data"]["start_time"] is not None
|
|
|
|
|
|
def test_export_dataset_metadata_enabled_by_config(
|
|
ray_start_cluster_with_export_api_config,
|
|
dummy_dataset_topology,
|
|
dummy_dataset_topology_expected_output,
|
|
):
|
|
_test_dataset_metadata_export(
|
|
dummy_dataset_topology, dummy_dataset_topology_expected_output
|
|
)
|
|
|
|
|
|
def test_export_dataset_metadata(
|
|
ray_start_cluster_with_export_api_write,
|
|
dummy_dataset_topology,
|
|
dummy_dataset_topology_expected_output,
|
|
):
|
|
_test_dataset_metadata_export(
|
|
dummy_dataset_topology, dummy_dataset_topology_expected_output
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"expected_logical_op_args",
|
|
[
|
|
{
|
|
"fn_args": [1],
|
|
"fn_constructor_kwargs": [2],
|
|
"fn_kwargs": {"a": 3},
|
|
"fn_constructor_args": {"b": 4},
|
|
"compute": ray.data.ActorPoolStrategy(max_tasks_in_flight_per_actor=2),
|
|
},
|
|
],
|
|
)
|
|
def test_logical_op_args(
|
|
ray_start_cluster_with_export_api_write, expected_logical_op_args
|
|
):
|
|
class Udf:
|
|
def __init__(self, a, b):
|
|
self.a = a
|
|
self.b = b
|
|
|
|
def __call__(self, x):
|
|
return x
|
|
|
|
ds = ray.data.range(1).map_batches(
|
|
Udf,
|
|
**expected_logical_op_args,
|
|
)
|
|
dag = ds._logical_plan.dag
|
|
args = dag._get_args()
|
|
assert len(args) > 0, "Export args should not be empty"
|
|
for k, v in expected_logical_op_args.items():
|
|
k = f"_{k}"
|
|
assert k in args, f"Export args should contain key '{k}'"
|
|
assert (
|
|
args[k] == v
|
|
), f"Export args for key '{k}' should match expected value {v}, found {args[k]}"
|
|
|
|
|
|
def test_export_multiple_datasets(
|
|
ray_start_cluster_with_export_api_write,
|
|
dummy_dataset_topology,
|
|
dummy_dataset_topology_expected_output,
|
|
):
|
|
"""Test that multiple datasets can be exported when export API is enabled."""
|
|
stats_actor = get_or_create_stats_actor()
|
|
|
|
# Create a second dataset structure that's different from the dummy one
|
|
second_topology = Topology(
|
|
operators=[
|
|
Operator(
|
|
name="Input",
|
|
id="Input_0",
|
|
uuid="second_uuid_0",
|
|
input_dependencies=[],
|
|
sub_stages=[],
|
|
execution_start_time=1.0,
|
|
execution_end_time=1.0,
|
|
state="FINISHED",
|
|
),
|
|
Operator(
|
|
name="ReadRange->Map(<lambda>)",
|
|
id="ReadRange->Map(<lambda>)_1",
|
|
uuid="second_uuid_1",
|
|
input_dependencies=["Input_0"],
|
|
sub_stages=[],
|
|
execution_start_time=2.0,
|
|
execution_end_time=0.0,
|
|
state="RUNNING",
|
|
),
|
|
],
|
|
)
|
|
|
|
# Dataset IDs for each dataset
|
|
first_dataset_id = "first_dataset"
|
|
second_dataset_id = "second_dataset"
|
|
|
|
# Register the first dataset
|
|
ray.get(
|
|
stats_actor.register_dataset.remote(
|
|
dataset_tag=first_dataset_id,
|
|
operator_tags=["ReadRange->Map(<lambda>)->Filter(<lambda>)"],
|
|
topology=dummy_dataset_topology,
|
|
job_id=STUB_JOB_ID,
|
|
data_context=DataContextMetadata.from_data_context(
|
|
DataContext.get_current()
|
|
),
|
|
)
|
|
)
|
|
|
|
# Register the second dataset
|
|
ray.get(
|
|
stats_actor.register_dataset.remote(
|
|
dataset_tag=second_dataset_id,
|
|
operator_tags=["ReadRange->Map(<lambda>)"],
|
|
topology=second_topology,
|
|
job_id=STUB_JOB_ID,
|
|
data_context=DataContextMetadata.from_data_context(
|
|
DataContext.get_current()
|
|
),
|
|
)
|
|
)
|
|
|
|
# Check that export files were created with both datasets
|
|
data = _get_exported_data()
|
|
assert len(data) == 2, f"Expected 2 exported datasets, got {len(data)}"
|
|
|
|
# Create a map of dataset IDs to their exported data for easier verification
|
|
datasets_by_id = {entry["event_data"]["dataset_id"]: entry for entry in data}
|
|
|
|
# Verify first dataset
|
|
assert (
|
|
first_dataset_id in datasets_by_id
|
|
), f"First dataset {first_dataset_id} not found in exported data"
|
|
first_entry = datasets_by_id[first_dataset_id]
|
|
assert first_entry["source_type"] == "EXPORT_DATASET_METADATA"
|
|
assert (
|
|
first_entry["event_data"]["topology"] == dummy_dataset_topology_expected_output
|
|
)
|
|
assert first_entry["event_data"]["job_id"] == STUB_JOB_ID
|
|
assert first_entry["event_data"]["start_time"] is not None
|
|
|
|
# Verify second dataset
|
|
assert (
|
|
second_dataset_id in datasets_by_id
|
|
), f"Second dataset {second_dataset_id} not found in exported data"
|
|
second_entry = datasets_by_id[second_dataset_id]
|
|
assert second_entry["source_type"] == "EXPORT_DATASET_METADATA"
|
|
assert second_entry["event_data"]["topology"] == asdict(second_topology)
|
|
assert second_entry["event_data"]["job_id"] == STUB_JOB_ID
|
|
assert second_entry["event_data"]["start_time"] is not None
|
|
|
|
|
|
class UnserializableObject:
|
|
"""A test class that can't be JSON serialized or converted to string easily."""
|
|
|
|
def __str__(self):
|
|
raise ValueError("Cannot convert to string")
|
|
|
|
def __repr__(self):
|
|
raise ValueError("Cannot convert to repr")
|
|
|
|
|
|
class BasicObject:
|
|
"""A test class that can be converted to string."""
|
|
|
|
def __init__(self, value):
|
|
self.value = value
|
|
|
|
def __str__(self):
|
|
return f"BasicObject({self.value})"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"input_obj,expected_output,truncate_length",
|
|
[
|
|
# Basic types - should return as strings
|
|
(42, "42", 100),
|
|
(3.14, "3.14", 100),
|
|
(True, "True", 100),
|
|
(False, "False", 100),
|
|
(None, "None", 100),
|
|
# Strings - short strings return as-is
|
|
("hello", "hello", 100),
|
|
# Strings - long strings get truncated
|
|
("a" * 150, "a" * 100 + "...", 100),
|
|
("hello world", "hello...", 5),
|
|
# Mappings - should recursively sanitize values
|
|
({"key": "value"}, {"key": "value"}, 100),
|
|
({"long_key": "a" * 150}, {"long_key": "a" * 100 + "..."}, 100),
|
|
({"nested": {"inner": "value"}}, {"nested": {"inner": "value"}}, 100),
|
|
# Sequences - should recursively sanitize elements (convert to strings)
|
|
([1, 2, 3], ["1", "2", "3"], 100),
|
|
(["short", "a" * 150], ["short", "a" * 100 + "..."], 100),
|
|
# Complex nested structures
|
|
(
|
|
{"list": [1, "a" * 150], "dict": {"key": "a" * 150}},
|
|
{"list": ["1", "a" * 100 + "..."], "dict": {"key": "a" * 100 + "..."}},
|
|
100,
|
|
),
|
|
# Objects that can be converted to string
|
|
(BasicObject("test"), "BasicObject(test)", 100), # Falls back to str()
|
|
# Sets can be converted to Lists of strings
|
|
({1, 2, 3}, ["1", "2", "3"], 100),
|
|
((1, 2, 3), ["1", "2", "3"], 100),
|
|
# Objects that can't be serialized or stringified
|
|
(UnserializableObject(), f"{UNKNOWN}: {UnserializableObject.__name__}", 100),
|
|
# Empty containers
|
|
({}, {}, 100),
|
|
([], [], 100),
|
|
# Mixed type sequences - all converted to strings
|
|
(
|
|
[1, "hello", {"key": "value"}, None],
|
|
["1", "hello", {"key": "value"}, "None"],
|
|
100,
|
|
),
|
|
# Bytearrays/bytes - should be converted to lists of string representations
|
|
(bytearray(b"hello"), ["104", "101", "108", "108", "111"], 100),
|
|
(bytearray([1, 2, 3, 4, 5]), ["1", "2", "3", "4", "5"], 100),
|
|
(bytes(b"test"), ["116", "101", "115", "116"], 100),
|
|
# Dataclass
|
|
(
|
|
TestDataclass(),
|
|
{
|
|
"list_field": ["1", "2", "3"],
|
|
"dict_field": {"1": "2", "3": "4"}, # key should be strings
|
|
"string_field": "test",
|
|
"int_field": "1",
|
|
"float_field": "1.0",
|
|
"set_field": [
|
|
"1",
|
|
"2",
|
|
"3",
|
|
], # sets will be converted to Lists of strings
|
|
"tuple_field": [
|
|
"1",
|
|
"2",
|
|
"3",
|
|
], # tuples will be converted to Lists of strings
|
|
"bool_field": "True",
|
|
"none_field": "None",
|
|
},
|
|
100,
|
|
),
|
|
# Test sequence truncation - list longer than truncate_length gets truncated
|
|
(
|
|
[1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
|
|
["1", "2", "3", "..."], # Only first 3 elements after truncation + ...
|
|
3,
|
|
),
|
|
],
|
|
)
|
|
def test_sanitize_for_struct(input_obj, expected_output, truncate_length):
|
|
"""Test sanitize_for_struct with various input types and truncation lengths."""
|
|
result = sanitize_for_struct(input_obj, truncate_length)
|
|
assert result == expected_output, f"Expected {expected_output}, got {result}"
|
|
|
|
|
|
def test_update_dataset_metadata_state(
|
|
ray_start_cluster_with_export_api_write, dummy_dataset_topology
|
|
):
|
|
"""Test dataset state update at the export API"""
|
|
stats_actor = get_or_create_stats_actor()
|
|
# Register dataset
|
|
ray.get(
|
|
stats_actor.register_dataset.remote(
|
|
job_id=STUB_JOB_ID,
|
|
dataset_tag=STUB_DATASET_ID,
|
|
operator_tags=["Input_0", "ReadRange->Map(<lambda>)->Filter(<lambda>)_1"],
|
|
topology=dummy_dataset_topology,
|
|
data_context=DataContextMetadata.from_data_context(
|
|
DataContext.get_current()
|
|
),
|
|
)
|
|
)
|
|
# Check that export files were created as expected
|
|
data = _get_exported_data()
|
|
assert len(data) == 1
|
|
assert data[0]["event_data"]["state"] == DatasetState.PENDING.name
|
|
|
|
# Test update state to RUNNING
|
|
ray.get(
|
|
stats_actor.update_dataset_metadata_state.remote(
|
|
dataset_id=STUB_DATASET_ID, new_state=DatasetState.RUNNING.name
|
|
)
|
|
)
|
|
data = _get_exported_data()
|
|
assert len(data) == 2
|
|
assert data[1]["event_data"]["state"] == DatasetState.RUNNING.name
|
|
assert data[1]["event_data"]["execution_start_time"] > 0
|
|
|
|
# Test update to FINISHED
|
|
ray.get(
|
|
stats_actor.update_dataset_metadata_state.remote(
|
|
dataset_id=STUB_DATASET_ID, new_state=DatasetState.FINISHED.name
|
|
)
|
|
)
|
|
data = _get_exported_data()
|
|
assert len(data) == 3
|
|
assert data[2]["event_data"]["state"] == DatasetState.FINISHED.name
|
|
assert data[2]["event_data"]["execution_end_time"] > 0
|
|
assert (
|
|
data[2]["event_data"]["topology"]["operators"][1]["state"]
|
|
== DatasetState.FINISHED.name
|
|
)
|
|
assert data[2]["event_data"]["topology"]["operators"][1]["execution_end_time"] > 0
|
|
|
|
|
|
def test_update_dataset_metadata_operator_states(
|
|
ray_start_cluster_with_export_api_write, dummy_dataset_topology
|
|
):
|
|
stats_actor = get_or_create_stats_actor()
|
|
# Register dataset
|
|
ray.get(
|
|
stats_actor.register_dataset.remote(
|
|
dataset_tag=STUB_DATASET_ID,
|
|
operator_tags=["Input_0", "ReadRange->Map(<lambda>)->Filter(<lambda>)_1"],
|
|
topology=dummy_dataset_topology,
|
|
job_id=STUB_JOB_ID,
|
|
data_context=DataContextMetadata.from_data_context(
|
|
DataContext.get_current()
|
|
),
|
|
)
|
|
)
|
|
data = _get_exported_data()
|
|
assert len(data) == 1
|
|
assert (
|
|
data[0]["event_data"]["topology"]["operators"][1]["state"]
|
|
== DatasetState.RUNNING.name
|
|
)
|
|
|
|
# Test update to FINISHED
|
|
operator_id = "ReadRange->Map(<lambda>)->Filter(<lambda>)_1"
|
|
ray.get(
|
|
stats_actor.update_dataset_metadata_operator_states.remote(
|
|
dataset_id=STUB_DATASET_ID,
|
|
operator_states={operator_id: DatasetState.FINISHED.name},
|
|
)
|
|
)
|
|
data = _get_exported_data()
|
|
assert len(data) == 2
|
|
assert (
|
|
data[1]["event_data"]["topology"]["operators"][1]["state"]
|
|
== DatasetState.FINISHED.name
|
|
)
|
|
assert data[1]["event_data"]["topology"]["operators"][1]["execution_end_time"] > 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|