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
@@ -39,6 +40,7 @@
 MULTIMODAL_CONTENT_TYPES = frozenset(
     {*IMAGE_CONTENT_TYPES, *VIDEO_CONTENT_TYPES, *AUDIO_CONTENT_TYPES}
 )
+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>"
@@ -888,6 +890,16 @@ 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 image source."""
+    image = resolve_to_image(source)
+    try:
+        data_url = image_to_data_url(image)
+    finally:
+        image.close()
+    return data_url
+
+
 def extract_input_image_sources_from_responses_messages(
     messages: Any,
 ) -> list[str | Image.Image]:
@@ -1033,13 +1045,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 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.
@@ -1053,6 +1065,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.
     """
+    targets_by_source: dict[str, list[tuple[dict, str]]] = {}
+
     for example in nemo_gym_examples:
         input_items = example.get("responses_create_params", {}).get("input", [])
         if not isinstance(input_items, list):
@@ -1082,7 +1096,27 @@ def encode_images_in_examples(nemo_gym_examples: list[dict]) -> list[dict]:
                     continue
                 if url.startswith(("http://", "https://", "data:")):
                     continue
-                part[media_key] = image_to_data_url(resolve_to_image(url))
+                targets_by_source.setdefault(url, []).append((part, media_key))
+
+    sources = list(targets_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, targets in targets_by_source.items():
+            data_url = encoded_by_source[source]
+            for part, media_key in targets:
+                part[media_key] = 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,
@@ -35,6 +39,14 @@ def _write_png(tmp_path, name: str, size: tuple[int, int]) -> str:
     return str(path)
 
 
+class _CloseTrackingImage:
+    def __init__(self) -> None:
+        self.closed: bool = False
+
+    def close(self) -> None:
+        self.closed = True
+
+
 def test_image_to_data_url_round_trips_through_resolve_to_image():
     url = image_to_data_url(Image.new("RGB", (4, 3)))
     assert url.startswith("data:image/png;base64,")
@@ -86,6 +98,100 @@ 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_a_bounded_thread_pool(
+    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": "image", "image": first},
+            {"type": "image_url", "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 == [multimodal_utils.NEMO_GYM_IMAGE_ENCODE_MAX_WORKERS]
+    for example in examples:
+        parts = example["responses_create_params"]["input"][0]["content"]
+        assert parts[0]["image_url"] == expected[first]
+        assert parts[1]["image"] == expected[first]
+        assert parts[2]["url"] == expected[second]
+
+
+def test_encode_images_does_not_partially_mutate_on_error(tmp_path):
+    existing = _write_png(tmp_path, "existing.png", (2, 3))
+    missing = str(tmp_path / "missing.png")
+    examples = [
+        _example(
+            {"type": "input_image", "image_url": existing},
+            {"type": "input_image", "image_url": missing},
+        )
+    ]
+
+    with pytest.raises(FileNotFoundError):
+        encode_images_in_examples(examples)
+
+    parts = examples[0]["responses_create_params"]["input"][0]["content"]
+    assert parts[0]["image_url"] == existing
+    assert parts[1]["image_url"] == missing
+
+
+def test_encode_single_image_source_closes_image_on_success(monkeypatch):
+    image = _CloseTrackingImage()
+    monkeypatch.setattr(multimodal_utils, "resolve_to_image", lambda _: image)
+    monkeypatch.setattr(
+        multimodal_utils,
+        "image_to_data_url",
+        lambda _: "data:image/png;base64,AA",
+    )
+
+    assert (
+        multimodal_utils._encode_single_image_source("image.png")
+        == "data:image/png;base64,AA"
+    )
+    assert image.closed
+
+
+def test_encode_single_image_source_closes_image_on_error(monkeypatch):
+    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 = [