diff --git a/nemo_rl/data/multimodal_utils.py b/nemo_rl/data/multimodal_utils.py index dda616b15c9..2a746a36097 100644 --- a/nemo_rl/data/multimodal_utils.py +++ b/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 " @@ -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 diff --git a/tests/unit/data/test_multimodal_image_encoding.py b/tests/unit/data/test_multimodal_image_encoding.py index e36c64135e0..fda52a35ca6 100644 --- a/tests/unit/data/test_multimodal_image_encoding.py +++ b/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 = [