wehub-resource-sync
66 lines
2.0 KiB
Python
66 lines
2.0 KiB
Python
from sklearn.datasets import load_breast_cancer
|
|
|
|
from ray import tune
|
|
from ray.data import Dataset, Datasource, ReadTask, read_datasource
|
|
from ray.data.block import BlockMetadata
|
|
from ray.tune.impl.utils import execute_dataset
|
|
|
|
|
|
# TODO(xwjiang): Enable this when Clark's out-of-band-serialization is landed.
|
|
class TestDatasource(Datasource):
|
|
def prepare_read(self, parallelism: int, **read_args):
|
|
import pyarrow as pa
|
|
|
|
def load_data():
|
|
data_raw = load_breast_cancer(as_frame=True)
|
|
dataset_df = data_raw["data"]
|
|
dataset_df["target"] = data_raw["target"]
|
|
return [pa.Table.from_pandas(dataset_df)]
|
|
|
|
meta = BlockMetadata(
|
|
num_rows=None,
|
|
size_bytes=None,
|
|
input_files=None,
|
|
exec_stats=None,
|
|
)
|
|
return [ReadTask(load_data, meta)]
|
|
|
|
|
|
def gen_dataset_func() -> Dataset:
|
|
test_datasource = TestDatasource()
|
|
return read_datasource(test_datasource)
|
|
|
|
|
|
def test_grid_search():
|
|
ds1 = gen_dataset_func().lazy().map(lambda x: x)
|
|
ds2 = gen_dataset_func().lazy().map(lambda x: x)
|
|
assert not ds1._has_computed_output()
|
|
assert not ds2._has_computed_output()
|
|
param_space = {"train_dataset": tune.grid_search([ds1, ds2])}
|
|
execute_dataset(param_space)
|
|
executed_ds = param_space["train_dataset"]["grid_search"]
|
|
assert len(executed_ds) == 2
|
|
assert executed_ds[0]._has_computed_output()
|
|
assert executed_ds[1]._has_computed_output()
|
|
|
|
|
|
def test_choice():
|
|
ds1 = gen_dataset_func().lazy().map(lambda x: x)
|
|
ds2 = gen_dataset_func().lazy().map(lambda x: x)
|
|
assert not ds1._has_computed_output()
|
|
assert not ds2._has_computed_output()
|
|
param_space = {"train_dataset": tune.choice([ds1, ds2])}
|
|
execute_dataset(param_space)
|
|
executed_ds = param_space["train_dataset"].categories
|
|
assert len(executed_ds) == 2
|
|
assert executed_ds[0]._has_computed_output()
|
|
assert executed_ds[1]._has_computed_output()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
import pytest
|
|
|
|
sys.exit(pytest.main(["-v", "-x", __file__]))
|