chore: import upstream snapshot with attribution

python/ray/data/tests/test_logging_dataset.py

@@ -0,0 +1,128 @@
+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__]))