Repository: ray-project/ray
File: python/ray/data/tests/test_consumption.py (Python)

Code before:
            [local_path + "/test1.parquet", data_path + "/test2.parquet"]
        ).materialize()
    with pytest.raises(ValueError):
        ds = ray.data.read_parquet(
            [local_path + "/test1.parquet", "example://iris.parquet"]
        ).materialize()
    with pytest.raises(ValueError):
        ds = ray.data.read_parquet(
            ["example://iris.parquet", local_path + "/test1.parquet"]
        ).materialize()


def test_validate_head_node_resources_zero_head_cpu(ray_start_cluster):
    from ray.data._internal.util import (
        _validate_head_node_resources_for_local_scheduling,
    )

    cluster = ray_start_cluster
    cluster.add_node(num_cpus=0)
    cluster.wait_for_nodes()

    ray.shutdown()
    ray.init(address=cluster.address)

    cluster.add_node(num_cpus=1)
    cluster.wait_for_nodes()

    with pytest.raises(ValueError, match=r"head node doesn't have enough resources"):
        _validate_head_node_resources_for_local_scheduling({}, op_name="read local")


class FlakyCSVDatasource(CSVDatasource):
    def __init__(self, paths, **csv_datasource_kwargs):
        super().__init__(paths, **csv_datasource_kwargs)

        self.counter = Counter.remote()

    def _read_stream(self, f: "pa.NativeFile", path: str):
        count = self.counter.increment.remote()
        if ray.get(count) == 1:
            raise ValueError("oops")
        else:
            for block in CSVDatasource._read_stream(self, f, path):
                yield block


class FlakyCSVDatasink(CSVDatasink):
    def __init__(self, path, **csv_datasink_kwargs):
        super().__init__(path, **csv_datasink_kwargs)

        self.counter = Counter.remote()

Review comment (refactor):
Nit: Is this logic necessary for the assertion below? If not, would it make sense to remove it to simplify the test?

Code after:
    # Mixing paths of local and non-local scheme is invalid.
    with pytest.raises(ValueError):
        ds = ray.data.read_parquet(
            [local_path + "/test1.parquet", data_path + "/test2.parquet"]
        ).materialize()
    with pytest.raises(ValueError):
        ds = ray.data.read_parquet(
            [local_path + "/test1.parquet", "example://iris.parquet"]
        ).materialize()
    with pytest.raises(ValueError):
        ds = ray.data.read_parquet(
            ["example://iris.parquet", local_path + "/test1.parquet"]
        ).materialize()


def test_validate_head_node_resources_zero_head_cpu(ray_start_cluster):

    cluster = ray_start_cluster
    cluster.add_node(num_cpus=0)
    cluster.wait_for_nodes()

    ray.shutdown()
    ray.init(address=cluster.address)

    with pytest.raises(ValueError, match=r"head node doesn't have enough resources"):
        _validate_head_node_resources_for_local_scheduling(
            {}, op_description="read local"
        )


class FlakyCSVDatasource(CSVDatasource):
    def __init__(self, paths, **csv_datasource_kwargs):
        super().__init__(paths, **csv_datasource_kwargs)

        self.counter = Counter.remote()

    def _read_stream(self, f: "pa.NativeFile", path: str):
        count = self.counter.increment.remote()
        if ray.get(count) == 1:
            raise ValueError("oops")
        else:
            for block in CSVDatasource._read_stream(self, f, path):
                yield block


class FlakyCSVDatasink(CSVDatasink):
    def __init__(self, path, **csv_datasink_kwargs):
        super().__init__(path, **csv_datasink_kwargs)

        self.counter = Counter.remote()