[Data] Remove Ray Data read parallelism from all libraries and tests (#…

…43692)

This PR is to remove Ray Data read `parallelism`, replace with `override_num_blocks` from all libraries and unit tests. The PR change on Ray Data is #43113 .

Signed-off-by: Cheng Su <[email protected]>

dashboard/modules/data/tests/test_data_head.py

@@ -37,7 +37,7 @@
 )
 def test_get_datasets():
     ray.init()
-    ds = ray.data.range(100, parallelism=20).map_batches(lambda x: x)
+    ds = ray.data.range(100, override_num_blocks=20).map_batches(lambda x: x)
     ds._set_name("data_head_test")
     ds.materialize()
 

python/ray/air/tests/test_new_dataset_config.py

@@ -239,15 +239,15 @@ def train_loop_per_worker():
 
 
 def test_per_epoch_preprocessing(ray_start_4_cpus):
-    ds = ray.data.range(100, parallelism=100).randomize_block_order()
+    ds = ray.data.range(100, override_num_blocks=100).randomize_block_order()
     test = TestRandom(2, True, datasets={"train": ds})
     test.fit()
 
-    ds = ray.data.range(100, parallelism=100).random_shuffle()
+    ds = ray.data.range(100, override_num_blocks=100).random_shuffle()
     test = TestRandom(2, True, datasets={"train": ds})
     test.fit()
 
-    ds = ray.data.range(100, parallelism=100).map(
+    ds = ray.data.range(100, override_num_blocks=100).map(
         lambda x: {"id": x["id"] * random.random()}
     )
     test = TestRandom(2, True, datasets={"train": ds})
@@ -257,7 +257,7 @@ def test_per_epoch_preprocessing(ray_start_4_cpus):
 def test_materialized_preprocessing(ray_start_4_cpus):
     # TODO(ekl) we should test all these configs with splitting enabled, but this
     # requires implementing deterministic streaming split.
-    ds = ray.data.range(100, parallelism=100).randomize_block_order()
+    ds = ray.data.range(100, override_num_blocks=100).randomize_block_order()
     ds = ds.materialize()
     test = TestRandom(
         2,
@@ -267,7 +267,7 @@ def test_materialized_preprocessing(ray_start_4_cpus):
     )
     test.fit()
 
-    ds = ray.data.range(100, parallelism=100).random_shuffle()
+    ds = ray.data.range(100, override_num_blocks=100).random_shuffle()
     ds = ds.materialize()
     test = TestRandom(
         2,
@@ -277,7 +277,7 @@ def test_materialized_preprocessing(ray_start_4_cpus):
     )
     test.fit()
 
-    ds = ray.data.range(100, parallelism=100).map(
+    ds = ray.data.range(100, override_num_blocks=100).map(
         lambda x: {"id": x["id"] * random.random()}
     )
     ds = ds.materialize()

python/ray/air/util/check_ingest.py

@@ -169,8 +169,8 @@ def make_local_dataset_iterator(
     args = parser.parse_args()
 
     # Generate a synthetic dataset of ~10GiB of float64 data. The dataset is sharded
-    # into 100 blocks (parallelism=100).
-    ds = ray.data.range_tensor(50000, shape=(80, 80, 4), parallelism=100)
+    # into 100 blocks (override_num_blocks=100).
+    ds = ray.data.range_tensor(50000, shape=(80, 80, 4), override_num_blocks=100)
 
     # An example preprocessing chain that just scales all values by 4.0 in two stages.
     ds = ds.map_batches(lambda df: df * 2, batch_format="pandas")

python/ray/data/dataset.py

@@ -4398,7 +4398,7 @@ def to_pandas_refs(self) -> List[ObjectRef["pandas.DataFrame"]]:
 
         Examples:
             >>> import ray
-            >>> ds = ray.data.range(10, parallelism=2)
+            >>> ds = ray.data.range(10, override_num_blocks=2)
             >>> refs = ds.to_pandas_refs()
             >>> len(refs)
             2
@@ -4426,7 +4426,7 @@ def to_numpy_refs(
 
         Examples:
             >>> import ray
-            >>> ds = ray.data.range(10, parallelism=2)
+            >>> ds = ray.data.range(10, override_num_blocks=2)
             >>> refs = ds.to_numpy_refs()
             >>> len(refs)
             2
@@ -4461,7 +4461,7 @@ def to_arrow_refs(self) -> List[ObjectRef["pyarrow.Table"]]:
 
         Examples:
             >>> import ray
-            >>> ds = ray.data.range(10, parallelism=2)
+            >>> ds = ray.data.range(10, override_num_blocks=2)
             >>> refs = ds.to_arrow_refs()
             >>> len(refs)
             2

python/ray/data/read_api.py

@@ -468,7 +468,7 @@ def read_mongo(
         ...     collection="my_collection",
         ...     pipeline=[{"$match": {"col2": {"$gte": 0, "$lt": 100}}}, {"$sort": "sort_field"}], # noqa: E501
         ...     schema=Schema({"col1": pa.string(), "col2": pa.int64()}),
-        ...     parallelism=10,
+        ...     override_num_blocks=10,
         ... )
 
     Args:
@@ -1955,10 +1955,10 @@ def read_sql(
         ``LIMIT`` and ``OFFSET`` to fetch a subset of the rows. However, for many
         databases, ``OFFSET`` is slow.
 
-        As a workaround, set ``parallelism=1`` to directly fetch all rows in a single
-        task. Note that this approach requires all result rows to fit in the memory of
-        single task. If the rows don't fit, your program may raise an out of memory
-        error.
+        As a workaround, set ``override_num_blocks=1`` to directly fetch all rows in a
+        single task. Note that this approach requires all result rows to fit in the
+        memory of single task. If the rows don't fit, your program may raise an out of
+        memory error.
 
     Examples:
 

python/ray/data/tests/test_all_to_all.py

@@ -20,8 +20,10 @@
 
 
 def test_zip(ray_start_regular_shared):
-    ds1 = ray.data.range(5, parallelism=5)
-    ds2 = ray.data.range(5, parallelism=5).map(column_udf("id", lambda x: x + 1))
+    ds1 = ray.data.range(5, override_num_blocks=5)
+    ds2 = ray.data.range(5, override_num_blocks=5).map(
+        column_udf("id", lambda x: x + 1)
+    )
     ds = ds1.zip(ds2)
     assert ds.schema().names == ["id", "id_1"]
     assert ds.take() == named_values(
@@ -39,8 +41,8 @@ def test_zip_different_num_blocks_combinations(
     ray_start_regular_shared, num_blocks1, num_blocks2
 ):
     n = 12
-    ds1 = ray.data.range(n, parallelism=num_blocks1)
-    ds2 = ray.data.range(n, parallelism=num_blocks2).map(
+    ds1 = ray.data.range(n, override_num_blocks=num_blocks1)
+    ds2 = ray.data.range(n, override_num_blocks=num_blocks2).map(
         column_udf("id", lambda x: x + 1)
     )
     ds = ds1.zip(ds2)
@@ -70,11 +72,11 @@ def test_zip_different_num_blocks_split_smallest(
     num_blocks1 = 4
     num_blocks2 = 2
     ds1 = ray.data.from_items(
-        [{str(i): i for i in range(num_cols1)}] * n, parallelism=num_blocks1
+        [{str(i): i for i in range(num_cols1)}] * n, override_num_blocks=num_blocks1
     )
     ds2 = ray.data.from_items(
         [{str(i): i for i in range(num_cols1, num_cols1 + num_cols2)}] * n,
-        parallelism=num_blocks2,
+        override_num_blocks=num_blocks2,
     )
     ds = ds1.zip(ds2).materialize()
     num_blocks = ds._plan._snapshot_blocks.executed_num_blocks()
@@ -129,8 +131,8 @@ def foo(x):
 
     num_items = 10
     items = list(range(num_items))
-    ds1 = ray.data.from_items(items, parallelism=num_items)
-    ds2 = ray.data.from_items(items, parallelism=num_items)
+    ds1 = ray.data.from_items(items, override_num_blocks=num_items)
+    ds2 = ray.data.from_items(items, override_num_blocks=num_items)
     ds2 = ds2.map_batches(foo, batch_size=1)
     result = ds1.zip(ds2).take_all()
     assert result == named_values(
@@ -139,15 +141,15 @@ def foo(x):
 
 
 def test_empty_shuffle(ray_start_regular_shared):
-    ds = ray.data.range(100, parallelism=100)
+    ds = ray.data.range(100, override_num_blocks=100)
     ds = ds.filter(lambda x: x)
     ds = ds.map_batches(lambda x: x)
     ds = ds.random_shuffle()  # Would prev. crash with AssertionError: pyarrow.Table.
     ds.show()
 
 
 def test_repartition_shuffle(ray_start_regular_shared):
-    ds = ray.data.range(20, parallelism=10)
+    ds = ray.data.range(20, override_num_blocks=10)
     assert ds._plan.initial_num_blocks() == 10
     assert ds.sum() == 190
     assert ds._block_num_rows() == [2] * 10
@@ -162,13 +164,13 @@ def test_repartition_shuffle(ray_start_regular_shared):
     assert ds3.sum() == 190
     assert ds3._block_num_rows() == [2] * 10 + [0] * 10
 
-    large = ray.data.range(10000, parallelism=10)
+    large = ray.data.range(10000, override_num_blocks=10)
     large = large.repartition(20, shuffle=True)
     assert large._block_num_rows() == [500] * 20
 
 
 def test_repartition_noshuffle(ray_start_regular_shared):
-    ds = ray.data.range(20, parallelism=10)
+    ds = ray.data.range(20, override_num_blocks=10)
     assert ds._plan.initial_num_blocks() == 10
     assert ds.sum() == 190
     assert ds._block_num_rows() == [2] * 10
@@ -194,13 +196,13 @@ def test_repartition_noshuffle(ray_start_regular_shared):
     assert ds5._plan.initial_num_blocks() == 4
     assert ds5._block_num_rows() == [5, 6, 5, 6]
 
-    large = ray.data.range(10000, parallelism=10)
+    large = ray.data.range(10000, override_num_blocks=10)
     large = large.repartition(20)
     assert large._block_num_rows() == [500] * 20
 
 
 def test_repartition_shuffle_arrow(ray_start_regular_shared):
-    ds = ray.data.range(20, parallelism=10)
+    ds = ray.data.range(20, override_num_blocks=10)
     assert ds._plan.initial_num_blocks() == 10
     assert ds.count() == 20
     assert ds._block_num_rows() == [2] * 10
@@ -215,7 +217,7 @@ def test_repartition_shuffle_arrow(ray_start_regular_shared):
     assert ds3.count() == 20
     assert ds3._block_num_rows() == [2] * 10 + [0] * 10
 
-    large = ray.data.range(10000, parallelism=10)
+    large = ray.data.range(10000, override_num_blocks=10)
     large = large.repartition(20, shuffle=True)
     assert large._block_num_rows() == [500] * 20
 
@@ -1132,7 +1134,7 @@ def test_random_block_order(ray_start_regular_shared, restore_data_context):
     assert results == expected
 
     # Test LazyBlockList.randomize_block_order.
-    lazy_blocklist_ds = ray.data.range(12, parallelism=4)
+    lazy_blocklist_ds = ray.data.range(12, override_num_blocks=4)
     lazy_blocklist_ds = lazy_blocklist_ds.randomize_block_order(seed=0)
     lazy_blocklist_results = lazy_blocklist_ds.take()
     lazy_blocklist_expected = named_values("id", [6, 7, 8, 0, 1, 2, 3, 4, 5, 9, 10, 11])
@@ -1148,8 +1150,8 @@ def test_random_shuffle(shutdown_only, use_push_based_shuffle):
     r2 = ray.data.range(100).random_shuffle().take(999)
     assert r1 != r2, (r1, r2)
 
-    r1 = ray.data.range(100, parallelism=1).random_shuffle().take(999)
-    r2 = ray.data.range(100, parallelism=1).random_shuffle().take(999)
+    r1 = ray.data.range(100, override_num_blocks=1).random_shuffle().take(999)
+    r2 = ray.data.range(100, override_num_blocks=1).random_shuffle().take(999)
     assert r1 != r2, (r1, r2)
 
     assert (
@@ -1160,22 +1162,22 @@ def test_random_shuffle(shutdown_only, use_push_based_shuffle):
     r2 = ray.data.range(100).random_shuffle().repartition(1).take(999)
     assert r1 != r2, (r1, r2)
 
-    r0 = ray.data.range(100, parallelism=5).take(999)
-    r1 = ray.data.range(100, parallelism=5).random_shuffle(seed=0).take(999)
-    r2 = ray.data.range(100, parallelism=5).random_shuffle(seed=0).take(999)
-    r3 = ray.data.range(100, parallelism=5).random_shuffle(seed=12345).take(999)
+    r0 = ray.data.range(100, override_num_blocks=5).take(999)
+    r1 = ray.data.range(100, override_num_blocks=5).random_shuffle(seed=0).take(999)
+    r2 = ray.data.range(100, override_num_blocks=5).random_shuffle(seed=0).take(999)
+    r3 = ray.data.range(100, override_num_blocks=5).random_shuffle(seed=12345).take(999)
     assert r1 == r2, (r1, r2)
     assert r1 != r0, (r1, r0)
     assert r1 != r3, (r1, r3)
 
-    r0 = ray.data.range(100, parallelism=5).take(999)
-    r1 = ray.data.range(100, parallelism=5).random_shuffle(seed=0).take(999)
-    r2 = ray.data.range(100, parallelism=5).random_shuffle(seed=0).take(999)
+    r0 = ray.data.range(100, override_num_blocks=5).take(999)
+    r1 = ray.data.range(100, override_num_blocks=5).random_shuffle(seed=0).take(999)
+    r2 = ray.data.range(100, override_num_blocks=5).random_shuffle(seed=0).take(999)
     assert r1 == r2, (r1, r2)
     assert r1 != r0, (r1, r0)
 
     # Test move.
-    ds = ray.data.range(100, parallelism=2)
+    ds = ray.data.range(100, override_num_blocks=2)
     r1 = ds.random_shuffle().take(999)
     ds = ds.map(lambda x: x).take(999)
     r2 = ray.data.range(100).random_shuffle().take(999)
@@ -1193,7 +1195,7 @@ def test_random_shuffle_check_random(shutdown_only):
     num_files = 10
     num_rows = 100
     items = [i for i in range(num_files) for _ in range(num_rows)]
-    ds = ray.data.from_items(items, parallelism=num_files)
+    ds = ray.data.from_items(items, override_num_blocks=num_files)
     out = ds.random_shuffle().take(num_files * num_rows)
     for i in range(num_files):
         part = out[i * num_rows : (i + 1) * num_rows]
@@ -1220,7 +1222,7 @@ def test_random_shuffle_check_random(shutdown_only):
     num_files = 10
     num_rows = 100
     items = [j for i in range(num_files) for j in range(num_rows)]
-    ds = ray.data.from_items(items, parallelism=num_files)
+    ds = ray.data.from_items(items, override_num_blocks=num_files)
     out = ds.random_shuffle().take(num_files * num_rows)
     for i in range(num_files):
         part = out[i * num_rows : (i + 1) * num_rows]
@@ -1252,7 +1254,7 @@ def test_random_shuffle_with_custom_resource(ray_start_cluster, use_push_based_s
     # Run dataset in "bar" nodes.
     ds = ray.data.read_parquet(
         "example://parquet_images_mini",
-        parallelism=2,
+        override_num_blocks=2,
         ray_remote_args={"resources": {"bar": 1}},
     )
     ds = ds.random_shuffle(resources={"bar": 1}).materialize()
@@ -1279,7 +1281,7 @@ def get_node_id():
     node1_id = ray.get(get_node_id.options(resources={"bar:1": 1}).remote())
     node2_id = ray.get(get_node_id.options(resources={"bar:2": 1}).remote())
 
-    ds = ray.data.range(100, parallelism=2).random_shuffle()
+    ds = ray.data.range(100, override_num_blocks=2).random_shuffle()
     blocks = ds.get_internal_block_refs()
     ray.wait(blocks, num_returns=len(blocks), fetch_local=False)
     location_data = ray.experimental.get_object_locations(blocks)

python/ray/data/tests/test_auto_parallelism.py

@@ -128,14 +128,14 @@ def test_auto_parallelism_basic(shutdown_only):
     context = DataContext.get_current()
     context.min_parallelism = 1
     # Datasource bound.
-    ds = ray.data.range_tensor(5, shape=(100,), parallelism=-1)
+    ds = ray.data.range_tensor(5, shape=(100,), override_num_blocks=-1)
     assert ds._plan.initial_num_blocks() == 5, ds
     # CPU bound. TODO(ekl) we should fix range datasource to respect parallelism more
     # properly, currently it can go a little over.
-    ds = ray.data.range_tensor(10000, shape=(100,), parallelism=-1)
+    ds = ray.data.range_tensor(10000, shape=(100,), override_num_blocks=-1)
     assert ds._plan.initial_num_blocks() == 16, ds
     # Block size bound.
-    ds = ray.data.range_tensor(100000000, shape=(100,), parallelism=-1)
+    ds = ray.data.range_tensor(100000000, shape=(100,), override_num_blocks=-1)
     assert ds._plan.initial_num_blocks() >= 590, ds
     assert ds._plan.initial_num_blocks() <= 600, ds
 
@@ -147,7 +147,7 @@ def test_auto_parallelism_placement_group(shutdown_only):
     def run():
         context = DataContext.get_current()
         context.min_parallelism = 1
-        ds = ray.data.range_tensor(2000, shape=(100,), parallelism=-1)
+        ds = ray.data.range_tensor(2000, shape=(100,), override_num_blocks=-1)
         return ds._plan.initial_num_blocks()
 
     # 1/16 * 4 * 16 = 4

python/ray/data/tests/test_backpressure_e2e.py

@@ -59,7 +59,7 @@ def consume(batch):
 
     last_snapshot = get_initial_core_execution_metrics_snapshot()
 
-    ds = ray.data.range(NUM_ROWS_TOTAL, parallelism=NUM_TASKS)
+    ds = ray.data.range(NUM_ROWS_TOTAL, override_num_blocks=NUM_TASKS)
     ds = ds.map_batches(produce, batch_size=NUM_ROWS_PER_TASK)
     ds = ds.map_batches(consume, batch_size=None, num_cpus=0.9)
     # Check core execution metrics every 10 rows, because it's expensive.
@@ -106,7 +106,7 @@ def consumer(batch):
         del batch["data"]
         return batch
 
-    ds = ray.data.range(1, parallelism=1).materialize()
+    ds = ray.data.range(1, override_num_blocks=1).materialize()
     ds = ds.map_batches(producer, batch_size=None, num_cpus=producer_num_cpus)
     # Add a limit op in the middle, to test that ReservationOpResourceAllocator
     # will account limit op's resource usage to the previous producer map op.
@@ -210,7 +210,7 @@ def map_fn(batch):
         batch["data"] = [np.zeros(block_size, dtype=np.uint8)]
         return batch
 
-    ds = ray.data.range(num_blocks, parallelism=num_blocks)
+    ds = ray.data.range(num_blocks, override_num_blocks=num_blocks)
     ds = ds.map_batches(map_fn, batch_size=None, num_cpus=1)
     assert len(ds.take_all()) == num_blocks
 
@@ -265,7 +265,7 @@ def range_(i):
     ctx.execution_options.resource_limits.object_store_memory = 10e6
 
     # 10GiB dataset.
-    ds = ray.data.read_datasource(source, n=10000, parallelism=1000)
+    ds = ray.data.read_datasource(source, n=10000, override_num_blocks=1000)
     it = iter(ds.iter_batches(batch_size=None, prefetch_batches=0))
     next(it)
     time.sleep(3)
@@ -290,7 +290,7 @@ def __call__(self, df: np.ndarray):
             return {"id": np.random.randn(1, 20, 1024, 1024)}
 
     ctx = ray.init(object_store_memory=4e9)
-    ds = ray.data.range_tensor(20, shape=(3, 1024, 1024), parallelism=20)
+    ds = ray.data.range_tensor(20, shape=(3, 1024, 1024), override_num_blocks=20)
 
     pipe = ds.map_batches(
         TestFast,