chore: import upstream snapshot with attribution

python/ray/tune/impl/test_utils.py

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