perf(data): parallelize NeMo Gym image encoding

nemo_rl/data/multimodal_utils.py

@@ -19,6 +19,7 @@
 import uuid
 from collections import defaultdict
 from collections.abc import Sequence
+from concurrent.futures import ThreadPoolExecutor
 from copy import deepcopy
 from io import BytesIO
 from typing import Any, Optional, Union
@@ -36,6 +37,7 @@
 MULTIMODAL_CONTENT_TYPES = frozenset(
     {"input_image", "image", "image_url", "video", "audio"}
 )
+NEMO_GYM_IMAGE_ENCODE_MAX_WORKERS = 8
 
 # List of allowed placeholder strings for different media types in the dataset string
 # e.g. "This is an example of <image>"
@@ -877,6 +879,15 @@ def image_to_data_url(image: Image.Image, fmt: str = "PNG") -> str:
     return f"data:image/{fmt.lower()};base64,{encoded}"
 
 
+def _encode_single_image_source(source: str) -> str:
+    """Resolve and encode one local image source."""
+    image = resolve_to_image(source)
+    try:
+        return image_to_data_url(image)
+    finally:
+        image.close()
+
+
 def extract_input_image_sources_from_responses_messages(
     messages: Any,
 ) -> list[str | Image.Image]:
@@ -1022,13 +1033,13 @@ def attach_image_model_inputs_to_message(
 def encode_images_in_examples(nemo_gym_examples: list[dict]) -> list[dict]:
     """Replace local image paths in NeMo Gym examples with base64 data URLs.
 
-    Walks each example's ``responses_create_params.input[].content[]`` items
-    and rewrites any ``input_image`` part whose ``image_url`` is a local path
-    (or ``file://`` URL) into a base64 ``data:`` URL via
-    :func:`image_to_data_url`. Parts whose URL already starts with ``http://``,
-    ``https://``, or ``data:`` are left untouched. Malformed items (non-dict
-    entries, missing/empty URLs, non-list ``input``/``content``) are skipped
-    without raising.
+    Walks each example's ``responses_create_params.input[].content[]`` items,
+    collects local image references, encodes each unique source once using a
+    bounded thread pool, and rewrites every corresponding ``input_image`` part
+    with the resulting base64 ``data:`` URL. Parts whose URL already starts
+    with ``http://``, ``https://``, or ``data:`` are left untouched. Malformed
+    items (non-dict entries, missing/empty URLs, non-list ``input``/``content``)
+    are skipped without raising.
 
     The examples are mutated in place; the same list is also returned for
     convenience so callers can chain the call.
@@ -1042,6 +1053,8 @@ def encode_images_in_examples(nemo_gym_examples: list[dict]) -> list[dict]:
         The same ``nemo_gym_examples`` list, with local image references
         rewritten to base64 data URLs in place.
     """
+    parts_by_source: dict[str, list[dict]] = {}
+
     for example in nemo_gym_examples:
         input_items = example.get("responses_create_params", {}).get("input", [])
         if not isinstance(input_items, list):
@@ -1062,7 +1075,27 @@ def encode_images_in_examples(nemo_gym_examples: list[dict]) -> list[dict]:
                     continue
                 if url.startswith(("http://", "https://", "data:")):
                     continue
-                part["image_url"] = image_to_data_url(resolve_to_image(url))
+                parts_by_source.setdefault(url, []).append(part)
+
+    sources = list(parts_by_source)
+    if sources:
+        with ThreadPoolExecutor(
+            max_workers=NEMO_GYM_IMAGE_ENCODE_MAX_WORKERS
+        ) as executor:
+            encoded_by_source = dict(
+                zip(
+                    sources,
+                    executor.map(_encode_single_image_source, sources),
+                    strict=True,
+                )
+            )
+
+        # Keep payload mutation on the caller thread after worker-owned images
+        # have been closed and every unique source has been encoded.
+        for source, parts in parts_by_source.items():
+            data_url = encoded_by_source[source]
+            for part in parts:
+                part["image_url"] = data_url
     return nemo_gym_examples
 
 

tests/unit/data/test_multimodal_image_encoding.py

@@ -12,8 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from collections import Counter
+
+import pytest
 from PIL import Image
 
+import nemo_rl.data.multimodal_utils as multimodal_utils
 from nemo_rl.data.multimodal_utils import (
     encode_images_in_examples,
     image_to_data_url,
@@ -86,6 +90,71 @@ def test_encode_images_passes_through_http_and_data_urls():
     assert parts[2]["image_url"] == data_url
 
 
+def test_encode_images_deduplicates_sources_and_uses_eight_workers(
+    tmp_path, monkeypatch
+):
+    first = _write_png(tmp_path, "first.png", (2, 3))
+    second = _write_png(tmp_path, "second.png", (4, 5))
+    expected = {
+        source: multimodal_utils._encode_single_image_source(source)
+        for source in (first, second)
+    }
+    examples = [
+        _example(
+            {"type": "input_image", "image_url": first},
+            {"type": "input_image", "image_url": second},
+        )
+        for _ in range(16)
+    ]
+
+    resolve_calls = []
+    original_resolve = multimodal_utils.resolve_to_image
+
+    def tracking_resolve(source):
+        resolve_calls.append(source)
+        return original_resolve(source)
+
+    observed_max_workers = []
+    original_executor = multimodal_utils.ThreadPoolExecutor
+
+    def tracking_executor(*args, **kwargs):
+        observed_max_workers.append(kwargs.get("max_workers"))
+        return original_executor(*args, **kwargs)
+
+    monkeypatch.setattr(multimodal_utils, "resolve_to_image", tracking_resolve)
+    monkeypatch.setattr(multimodal_utils, "ThreadPoolExecutor", tracking_executor)
+
+    encode_images_in_examples(examples)
+
+    assert Counter(resolve_calls) == Counter({first: 1, second: 1})
+    assert observed_max_workers == [8]
+    for example in examples:
+        parts = example["responses_create_params"]["input"][0]["content"]
+        assert parts[0]["image_url"] == expected[first]
+        assert parts[1]["image_url"] == expected[second]
+
+
+def test_encode_single_image_source_closes_image_on_error(monkeypatch):
+    class CloseTrackingImage:
+        def __init__(self):
+            self.closed = False
+
+        def close(self):
+            self.closed = True
+
+    image = CloseTrackingImage()
+    monkeypatch.setattr(multimodal_utils, "resolve_to_image", lambda _: image)
+
+    def fail_to_encode(_):
+        raise RuntimeError("encoding failed")
+
+    monkeypatch.setattr(multimodal_utils, "image_to_data_url", fail_to_encode)
+
+    with pytest.raises(RuntimeError, match="encoding failed"):
+        multimodal_utils._encode_single_image_source("image.png")
+    assert image.closed
+
+
 def test_encode_images_is_a_noop_for_text_only_examples():
     examples = [_example({"type": "input_text", "text": "no images here"})]
     before = [