[workflow] Enhance dataset tests (ray-project#25876)

python/ray/workflow/tests/test_dataset.py

@@ -8,22 +8,35 @@
 
 @ray.remote
 def gen_dataset():
- # TODO(ekl) seems checkpointing hangs with nested refs of
- # LazyBlockList.
  return ray.data.range(1000).map(lambda x: x)
 
 
+@ray.remote
+def gen_dataset_1():
+ return ray.data.range(1000)
+
+
+@ray.remote
+def gen_dataset_2():
+ return ray.data.range_table(1000)
+
+
 @ray.remote
 def transform_dataset(in_data):
  return in_data.map(lambda x: x * 2)
 
 
+@ray.remote
+def transform_dataset_1(in_data):
+ return in_data.map(lambda r: {"v2": r["value"] * 2})
+
+
 @ray.remote
 def sum_dataset(ds):
  return ds.sum()
 
 
-def test_dataset(workflow_start_regular):
+def test_dataset(workflow_start_regular_shared):
  ds_ref = gen_dataset.bind()
  transformed_ref = transform_dataset.bind(ds_ref)
  output_ref = sum_dataset.bind(transformed_ref)
@@ -32,6 +45,24 @@ def test_dataset(workflow_start_regular):
  assert result == 2 * sum(range(1000))
 
 
+def test_dataset_1(workflow_start_regular_shared):
+ ds_ref = gen_dataset_1.bind()
+ transformed_ref = transform_dataset.bind(ds_ref)
+ output_ref = sum_dataset.bind(transformed_ref)
+
+ result = workflow.create(output_ref).run()
+ assert result == 2 * sum(range(1000))
+
+
+def test_dataset_2(workflow_start_regular_shared):
+ ds_ref = gen_dataset_2.bind()
+ transformed_ref = transform_dataset_1.bind(ds_ref)
+ output_ref = sum_dataset.bind(transformed_ref)
+
+ result = workflow.create(output_ref).run()
+ assert result == 2 * sum(range(1000))
+
+
 if __name__ == "__main__":
  import sys