wehub-resource-sync
129 lines
4.0 KiB
Python
129 lines
4.0 KiB
Python
import logging
|
|
import os
|
|
import time
|
|
|
|
import pytest
|
|
|
|
import ray
|
|
from ray.data._internal.logging import get_log_directory, reset_logging
|
|
from ray.tests.conftest import * # noqa
|
|
|
|
|
|
@pytest.fixture(name="reset_logging")
|
|
def reset_logging_fixture():
|
|
yield
|
|
reset_logging()
|
|
|
|
|
|
def _test_harness(filename_01: str, filename_02: str):
|
|
"""Test harness to check if the log files are created and contain the expected content."""
|
|
dataset_01_log_file = None
|
|
dataset_02_log_file = None
|
|
directory = get_log_directory()
|
|
for filename in os.listdir(directory):
|
|
if filename.startswith(f"ray-data-{filename_01}") and filename.endswith(".log"):
|
|
dataset_01_log_file = os.path.join(directory, filename)
|
|
if filename.startswith(f"ray-data-{filename_02}") and filename.endswith(".log"):
|
|
dataset_02_log_file = os.path.join(directory, filename)
|
|
if dataset_01_log_file and dataset_02_log_file:
|
|
break
|
|
|
|
assert dataset_01_log_file is not None
|
|
assert dataset_02_log_file is not None
|
|
with open(dataset_01_log_file, "r") as f:
|
|
log_01_contents = f.read()
|
|
with open(dataset_02_log_file, "r") as f:
|
|
log_02_contents = f.read()
|
|
dataset_01_id = (
|
|
os.path.basename(dataset_01_log_file)
|
|
.removeprefix("ray-data-")
|
|
.removesuffix(".log")
|
|
)
|
|
dataset_02_id = (
|
|
os.path.basename(dataset_02_log_file)
|
|
.removeprefix("ray-data-")
|
|
.removesuffix(".log")
|
|
)
|
|
|
|
assert dataset_01_id in log_01_contents
|
|
assert dataset_02_id in log_02_contents
|
|
assert (
|
|
f"Starting execution of Dataset {dataset_01_id}"
|
|
in log_01_contents + log_02_contents
|
|
)
|
|
assert (
|
|
f"Dataset {dataset_02_id} execution finished"
|
|
in log_01_contents + log_02_contents
|
|
)
|
|
return log_01_contents, log_02_contents
|
|
|
|
|
|
def test_dataset_logging_concurrent(ray_start_regular_shared, reset_logging):
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
|
|
def _short(x):
|
|
logger = logging.getLogger("ray.data")
|
|
logger.info("short function is running")
|
|
time.sleep(0.1)
|
|
return x
|
|
|
|
def _long(x):
|
|
logger = logging.getLogger("ray.data")
|
|
logger.info("long function is running")
|
|
time.sleep(1)
|
|
return x
|
|
|
|
ds01 = ray.data.range(1).map_batches(_short)
|
|
ds01.set_name("test_dataset_logging_concurrent_01")
|
|
|
|
ds02 = ray.data.range(1).map_batches(_long)
|
|
ds02.set_name("test_dataset_logging_concurrent_02")
|
|
|
|
with ThreadPoolExecutor() as executor:
|
|
executor.submit(ds01.materialize)
|
|
executor.submit(ds02.materialize)
|
|
|
|
log_01_contents, log_02_contents = _test_harness(
|
|
"test_dataset_logging_concurrent_01",
|
|
"test_dataset_logging_concurrent_02",
|
|
)
|
|
# for concurrent datasets, which dataset contains the worker logs is not
|
|
# deterministic, so we only check that the logs are present in the combined logs
|
|
assert "short function is running" in log_01_contents + log_02_contents
|
|
assert "long function is running" in log_01_contents + log_02_contents
|
|
|
|
|
|
def test_dataset_logging_sequential(ray_start_regular_shared, reset_logging):
|
|
def _short(x):
|
|
logger = logging.getLogger("ray.data")
|
|
logger.info("short function is running")
|
|
time.sleep(0.1)
|
|
return x
|
|
|
|
def _long(x):
|
|
logger = logging.getLogger("ray.data")
|
|
logger.info("long function is running")
|
|
time.sleep(1)
|
|
return x
|
|
|
|
ds01 = ray.data.range(1).map_batches(_short)
|
|
ds01.set_name("test_dataset_logging_sequential_01")
|
|
ds01.materialize()
|
|
|
|
ds02 = ray.data.range(1).map_batches(_long)
|
|
ds02.set_name("test_dataset_logging_sequential_02")
|
|
ds02.materialize()
|
|
|
|
log_01_contents, log_02_contents = _test_harness(
|
|
"test_dataset_logging_sequential_01",
|
|
"test_dataset_logging_sequential_02",
|
|
)
|
|
assert "short function is running" in log_01_contents
|
|
assert "long function is running" in log_02_contents
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|