Address comment and fix unit test

Signed-off-by: Cheng Su <[email protected]>
ray-project · Feb 26, 2024 · a2a2031 · a2a2031
1 parent 478ac16
commit a2a2031
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/python/ray/data/_internal/logical/rules/set_read_parallelism.py b/python/ray/data/_internal/logical/rules/set_read_parallelism.py
@@ -25,7 +25,9 @@ def compute_additional_split_factor(
  detected_parallelism, reason, _, _ = _autodetect_parallelism(
  parallelism, target_max_block_size, ctx, datasource_or_legacy_reader, mem_size
  )
- num_read_tasks = len(datasource_or_legacy_reader.get_read_tasks(detected_parallelism))
+ num_read_tasks = len(
+ datasource_or_legacy_reader.get_read_tasks(detected_parallelism)
+ )
  expected_block_size = None
  if mem_size:
  expected_block_size = mem_size / num_read_tasks
@@ -47,7 +49,7 @@ def compute_additional_split_factor(
 
  available_cpu_slots = ray_available_resources().get("CPU", 1)
  if (
- parallelism > 0
+ parallelism != -1
  and num_read_tasks >= available_cpu_slots * 4
  and num_read_tasks >= 5000
  ):

diff --git a/python/ray/data/tests/test_consumption.py b/python/ray/data/tests/test_consumption.py
@@ -1540,7 +1540,7 @@ def test_read_warning_large_parallelism(ray_start_regular, propagate_logs, caplo
  with caplog.at_level(logging.WARNING, logger="ray.data.read_api"):
  ray.data.range(5000, parallelism=5000).materialize()
  assert (
- "The requested parallelism of 5000 is "
+ "The requested number of read blocks of 5000 is "
  "more than 4x the number of available CPU slots in the cluster" in caplog.text
  ), caplog.text