[Data] Cherry-pick object detection example to 2.5.0 release (ray-pro…

…ject#35637) * [Data] Add batch inference object detection example (ray-project#35143) Signed-off-by: Hao Chen <chenh1024@gmail.com> * Fix object detection example test Signed-off-by: Hao Chen <chenh1024@gmail.com> --------- Signed-off-by: Hao Chen <chenh1024@gmail.com>
rkooo567 · May 23, 2023 · c0d657a · c0d657a
1 parent 896f834
commit c0d657a
Show file tree

Hide file tree

Showing 6 changed files with 1,090 additions and 2 deletions.
diff --git a/doc/source/_toc.yml b/doc/source/_toc.yml
@@ -102,6 +102,8 @@ parts:
               - file: data/examples/ocr_example
                 title: Scaling OCR with Ray Data
               - file: data/examples/random-access
+              - file: data/examples/batch_inference_object_detection
+                title: Object Detection Batch Inference with PyTorch
           - file: data/faq
           - file: data/api/api
           - file: data/glossary

diff --git a/doc/source/data/batch_inference.rst b/doc/source/data/batch_inference.rst
@@ -782,6 +782,14 @@ tutorials and examples:
 
             Batch Inference on NYC taxi data using Ray Data
 
+    .. grid-item-card::
+        :img-top: /images/ray_logo.png
+        :class-img-top: pt-2 w-75 d-block mx-auto fixed-height-img
+
+        .. button-ref:: /data/examples/batch_inference_object_detection
+
+            Object Detection Batch Inference with PyTorch
+
     .. grid-item-card::
         :img-top: /images/ray_logo.png
         :class-img-top: pt-2 w-75 d-block mx-auto fixed-height-img

diff --git a/doc/source/data/examples/BUILD b/doc/source/data/examples/BUILD
@@ -5,3 +5,11 @@ filegroup(
     srcs = glob(["*.ipynb"]),
     visibility = ["//doc:__subpackages__"]
 )
+
+py_test_run_all_notebooks(
+    size = "medium",
+    include = ["batch_inference_object_detection.ipynb"],
+    exclude = [],
+    data = ["//doc/source/data/examples:data_examples"],
+    tags = ["exclusive", "team:data", "gpu"],
+)
diff --git a/doc/source/data/examples/batch_inference_object_detection.ipynb b/doc/source/data/examples/batch_inference_object_detection.ipynb
diff --git a/doc/test_myst_doc.py b/doc/test_myst_doc.py
@@ -48,8 +48,13 @@ def postprocess_notebook(notebook):
     return notebook
 
 
-if __name__ == "__main__":
+DISPLAY_FUNCTION = """
+def display(*args, **kwargs):
+    print(*args, **kwargs)
+"""
 
+
+if __name__ == "__main__":
     args, remainder = parser.parse_known_args()
 
     path = Path(args.path)
@@ -66,6 +71,9 @@ def postprocess_notebook(notebook):
 
     name = ""
     with tempfile.NamedTemporaryFile("w", delete=False) as f:
+        # Define the display function, which is available in notebooks,
+        # but not in normal Python scripts.
+        f.write(DISPLAY_FUNCTION)
         jupytext.write(notebook, f, fmt="py:percent")
         name = f.name
 

diff --git a/python/ray/data/_internal/table_block.py b/python/ray/data/_internal/table_block.py
@@ -146,7 +146,10 @@ def _compact_if_needed(self) -> None:
         assert self._columns
         if self._uncompacted_size.size_bytes() < MAX_UNCOMPACTED_SIZE_BYTES:
             return
-        block = self._table_from_pydict(self._columns)
+        columns = {
+            key: convert_udf_returns_to_numpy(col) for key, col in self._columns.items()
+        }
+        block = self._table_from_pydict(columns)
         self.add_block(block)
         self._uncompacted_size = SizeEstimator()
         self._columns.clear()