wehub-resource-sync
99 lines
3.0 KiB
Python
99 lines
3.0 KiB
Python
import logging
|
|
import sys
|
|
from unittest.mock import patch
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pyarrow as pa
|
|
import pytest
|
|
|
|
import ray
|
|
from ray.data import Schema
|
|
from ray.data._internal.util import _check_pyarrow_version
|
|
from ray.data.tests.conftest import * # noqa
|
|
from ray.tests.conftest import * # noqa
|
|
|
|
|
|
def test_column_name_type_check(ray_start_regular_shared):
|
|
df = pd.DataFrame({"1": np.random.rand(10), "a": np.random.rand(10)})
|
|
ds = ray.data.from_pandas(df)
|
|
assert ds.schema() == Schema(pa.schema([("1", pa.float64()), ("a", pa.float64())]))
|
|
assert ds.count() == 10
|
|
|
|
df = pd.DataFrame({1: np.random.rand(10), "a": np.random.rand(10)})
|
|
with pytest.raises(ValueError):
|
|
ray.data.from_pandas(df)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
sys.version_info >= (3, 12), reason="TODO(scottjlee): Not working yet for py312"
|
|
)
|
|
def test_unsupported_pyarrow_versions_check(shutdown_only):
|
|
ray.shutdown()
|
|
|
|
# Test that unsupported pyarrow versions cause an error to be raised upon the
|
|
# initial pyarrow use.
|
|
# Pin numpy<2 since pyarrow 8.0.0 is incompatible with numpy>=2,
|
|
# which is now the default on the data CI images.
|
|
ray.init(runtime_env={"pip": ["numpy==1.26.4", "pyarrow==8.0.0"]})
|
|
|
|
@ray.remote
|
|
def should_error():
|
|
_check_pyarrow_version()
|
|
|
|
with pytest.raises(
|
|
Exception,
|
|
match=r".*Dataset requires pyarrow >= 17.0.0, but 8.0.0 is installed.*",
|
|
):
|
|
ray.get(should_error.remote())
|
|
|
|
|
|
class LoggerWarningCalled(Exception):
|
|
"""Custom exception used in test_warning_execute_with_no_cpu() and
|
|
test_nowarning_execute_with_cpu(). Raised when the `logger.warning` method
|
|
is called, so that we can kick out of `plan.execute()` by catching this Exception
|
|
and check logging was done properly."""
|
|
|
|
pass
|
|
|
|
|
|
def test_warning_execute_with_no_cpu(ray_start_cluster):
|
|
"""Tests Dataset._execute() to ensure a warning is logged
|
|
when no CPU resources are available."""
|
|
# Create one node with no CPUs to trigger the Dataset warning
|
|
ray.shutdown()
|
|
ray.init(ray_start_cluster.address)
|
|
cluster = ray_start_cluster
|
|
cluster.add_node(num_cpus=0)
|
|
|
|
try:
|
|
ds = ray.data.range(10)
|
|
ds = ds.map_batches(lambda x: x)
|
|
ds.take()
|
|
except Exception as e:
|
|
assert isinstance(e, ValueError)
|
|
assert "exceeds the execution limits ExecutionResources(cpu=0.0" in str(e)
|
|
|
|
|
|
def test_nowarning_execute_with_cpu(ray_start_cluster):
|
|
"""Tests Dataset._execute() to ensure no warning is logged
|
|
when there are available CPU resources."""
|
|
# Create one node with CPUs to avoid triggering the Dataset warning
|
|
ray.shutdown()
|
|
ray.init(ray_start_cluster.address)
|
|
|
|
logger = logging.getLogger("ray.data.dataset")
|
|
with patch.object(
|
|
logger,
|
|
"warning",
|
|
side_effect=LoggerWarningCalled,
|
|
) as mock_logger:
|
|
ds = ray.data.range(10)
|
|
ds = ds.map_batches(lambda x: x)
|
|
ds.take()
|
|
mock_logger.assert_not_called()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main(["-v", __file__]))
|