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

Code before:
        .materialize()
    )
    check_dataset_is_local(ds)

    # Write back to local scheme.
    output = os.path.join(local_path, "test_read_write_local_node")
    ds.write_parquet(output)
    assert "1 nodes used" in ds.stats(), ds.stats()
    ray.data.read_parquet(output).take_all() == ds.take_all()

    # 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_read_local_scheme_zero_head_resources(ray_start_cluster, tmp_path):
    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()

    csv_path = tmp_path / "data.csv"
    pd.DataFrame({"a": [1, 2]}).to_csv(csv_path, index=False)

    with pytest.raises(ValueError, match=r"local://.*head node"):
        ray.data.read_csv(f"local://{csv_path}").materialize()


def test_write_local_scheme_zero_head_resources(ray_start_cluster, tmp_path):
    cluster = ray_start_cluster
    cluster.add_node(num_cpus=0)
    cluster.wait_for_nodes()

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


Review comment (suggestion):
I think these tests test very similar things. What I recommend is just keep 1 of the tests and use `_validate_head_node_resources_for_local_scheduling` directly for testing (since it's not really testing read/write E2E), or you can `pytest.parameterize`? I prefer the 1st option since it's not high risk if missing, and can makes the testing a faster (since we only need one of the tests)

Code after:
    ds = (
        ray.data.read_parquet(local_path, override_num_blocks=1)
        .map(lambda x: x)
        .materialize()
    )
    check_dataset_is_local(ds)

    # Write back to local scheme.
    output = os.path.join(local_path, "test_read_write_local_node")
    ds.write_parquet(output)
    assert "1 nodes used" in ds.stats(), ds.stats()
    ray.data.read_parquet(output).take_all() == ds.take_all()

    # 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()