[AIR] Optimize computer vision example code (ray-project#38712)

The computer_vision.py example test previously took more than 15 minutes to run. With limit pushdown and the changes in this PR, this test now only takes ~6 minutes to run.

Signed-off-by: Balaji Veeramani <[email protected]>

doc/BUILD

@@ -200,22 +200,11 @@ py_test_run_all_subdirectory(
 py_test_run_all_subdirectory(
     size = "large",
     include = ["source/ray-air/doc_code/*.py"],
-    exclude = [
-        # Too large. Use a custom test below.
-        "source/ray-air/doc_code/computer_vision.py",
-    ],
+    exclude = [],
     extra_srcs = [],
     tags = ["exclusive", "team:ml"],
 )
 
-py_test(
-    size = "enormous",
-    name = "computer_vision_test",
-    main = "source/ray-air/doc_code/computer_vision.py",
-    srcs = ["source/ray-air/doc_code/computer_vision.py"],
-    tags = ["exclusive", "team:ml"],
-)
-
 # --------------------------------------------------------------------
 # Test all doc/source/train/doc_code code included in rst/md files.
 # --------------------------------------------------------------------

doc/source/ray-air/doc_code/computer_vision.py

@@ -97,7 +97,10 @@ def read_images():
     dataset = ray.data.read_images(root, partitioning=partitioning)
     # __read_images1_stop__
 
-    dataset = dataset.limit(32)
+    # The autodetected parallelism is low. As a result, blocks are large and we
+    # unnecessarily read more than a thousand images (even though we limit the dataset
+    # to 32 rows!) To avoid this issue, we manually set the parallelism.
+    dataset = ray.data.read_images(root, partitioning=partitioning, parallelism=1875)
 
     # __read_images2_start__
     from typing import Dict