From 2fa1db41cf857083d8e09b441501115aca00a070 Mon Sep 17 00:00:00 2001 From: Zhengyang Liu Date: Fri, 6 Nov 2026 12:16:21 +0800 Subject: [PATCH] fix(multimodal): preserve distinct image prompts --- agentlightning/_prompt_identity.py | 44 ++++++++++ agentlightning/server/routes/events.py | 30 ++++--- agentlightning/verl/agl_rollout_manager.py | 39 ++------- docs/80-example-multimodal-qa.md | 2 + tests/server/test_endpoints.py | 41 +++++++++ tests/verl/test_agl_rollout_manager.py | 99 +++++++++++++++++++++- 6 files changed, 211 insertions(+), 44 deletions(-) create mode 100644 agentlightning/_prompt_identity.py diff --git a/agentlightning/_prompt_identity.py b/agentlightning/_prompt_identity.py new file mode 100644 index 000000000..fb41dc928 --- /dev/null +++ b/agentlightning/_prompt_identity.py @@ -0,0 +1,44 @@ +# Copyright (c) Microsoft. All rights reserved. + +"""Prompt identity shared by event deduplication and image alignment.""" + +from __future__ import annotations + +import json +from typing import Any + +PromptKey = tuple[tuple[int, ...], tuple[str, ...]] + + +def extract_image_urls_from_messages(messages: Any) -> list[str]: + if not isinstance(messages, list): + return [] + image_urls: list[str] = [] + for message in messages: + if not isinstance(message, dict): + continue + content = message.get("content") + if isinstance(content, str): + try: + content = json.loads(content) # Some clients serialize content parts as JSON text. + except json.JSONDecodeError: + continue + if isinstance(content, list): + for content_part in content: + if not isinstance(content_part, dict) or content_part.get("type") != "image_url": + continue + image_url = content_part.get("image_url") + if isinstance(image_url, dict) and isinstance(image_url.get("url"), str): + image_urls.append(image_url["url"]) + return image_urls + + +def prompt_key(prompt_token_ids: Any, image_urls: list[str]) -> PromptKey | None: + """Include ordered image references: placeholder token IDs omit image content.""" + if ( + not isinstance(prompt_token_ids, list) + or not prompt_token_ids + or any(type(token_id) is not int for token_id in prompt_token_ids) + ): + return None + return tuple(prompt_token_ids), tuple(image_urls) diff --git a/agentlightning/server/routes/events.py b/agentlightning/server/routes/events.py index 35faab72a..8e2ca71ca 100644 --- a/agentlightning/server/routes/events.py +++ b/agentlightning/server/routes/events.py @@ -11,6 +11,7 @@ from fastapi import APIRouter, Query from fastapi.exceptions import HTTPException +from agentlightning._prompt_identity import PromptKey, extract_image_urls_from_messages, prompt_key from agentlightning.schemas import DEFAULT_ATTEMPT_ID, Event, EventCreate from agentlightning.server.store import _events, _rollouts @@ -28,10 +29,18 @@ def _drop_superseded_routed_experts(events: list[Event], event_type: str, data: current = _trim_model_request(data) prompt_ids = current["prompt_token_ids"] + request = data.get("request") + messages = request.get("messages") if isinstance(request, dict) else None + image_urls = extract_image_urls_from_messages(messages) for index in range(len(events) - 1, -1, -1): previous = events[index] if previous.event_type != "model_request": continue + previous_request = previous.data.get("request") + previous_messages = previous_request.get("messages") if isinstance(previous_request, dict) else None + if image_urls != extract_image_urls_from_messages(previous_messages): + # Equal placeholder IDs do not imply equal visual inputs or routes. + return previous_data = _trim_model_request(previous.data) previous_context = previous_data["prompt_token_ids"] + previous_data["response_token_ids"] same_prompt = prompt_ids == previous_data["prompt_token_ids"] @@ -209,23 +218,22 @@ def _without_routed_experts(event: Event) -> Event: return event.model_copy(update={"data": data}) -def _dedupe_model_requests_by_prompt_token_ids(events: list[Event]) -> list[Event]: - """Keep the last request for each valid, non-empty prompt-token key.""" - last_index_by_prompt: dict[tuple[int, ...], int] = {} +def _dedupe_model_requests_by_prompt(events: list[Event]) -> list[Event]: + """Keep the last raw request for each valid prompt-token and image key.""" + last_index_by_prompt: dict[PromptKey, int] = {} kept_indexes: set[int] = set() for index, event in enumerate(events): if event.event_type != "model_request": continue - prompt_token_ids = event.data.get("prompt_token_ids", []) - if ( - not isinstance(prompt_token_ids, list) - or not prompt_token_ids - or any(type(token_id) is not int for token_id in prompt_token_ids) - ): + prompt_token_ids = _trim_model_request(event.data)["prompt_token_ids"] + request = event.data.get("request") + messages = request.get("messages") if isinstance(request, dict) else None + key = prompt_key(prompt_token_ids, extract_image_urls_from_messages(messages)) + if key is None: # Skip deduplication, not the request, when no valid key exists. kept_indexes.add(index) continue - last_index_by_prompt[tuple(prompt_token_ids)] = index + last_index_by_prompt[key] = index kept_indexes.update(last_index_by_prompt.values()) return [event for index, event in enumerate(events) if event.event_type != "model_request" or index in kept_indexes] @@ -250,8 +258,8 @@ async def query_events( event_type=event_type, ) if format == "triplet": + events = _dedupe_model_requests_by_prompt(events) events = [_to_triplet_format(e) for e in events] - events = _dedupe_model_requests_by_prompt_token_ids(events) if not include_routed_experts: events = [_without_routed_experts(event) for event in events] return events diff --git a/agentlightning/verl/agl_rollout_manager.py b/agentlightning/verl/agl_rollout_manager.py index 792ecc821..2932b0acc 100644 --- a/agentlightning/verl/agl_rollout_manager.py +++ b/agentlightning/verl/agl_rollout_manager.py @@ -4,7 +4,6 @@ from __future__ import annotations -import json # [multimodal-patch] import time import traceback import uuid @@ -18,6 +17,8 @@ from httpx_retries import Retry, RetryTransport from pydantic import BaseModel, Field +from agentlightning._prompt_identity import PromptKey, prompt_key +from agentlightning._prompt_identity import extract_image_urls_from_messages as _extract_image_urls_from_messages from agentlightning.client import AgentLightningSyncClient from agentlightning.schemas import ( TERMINAL_STATES, @@ -141,33 +142,8 @@ def _without_routed_experts(event: Event) -> Event: return event.model_copy(update={"data": data}) -# [multimodal-patch] Extract image URLs from OpenAI-style chat messages, in order of -# appearance (ported from agent-lightning v0.3.0 TripletAdapter.extract_prompt_image_urls). -def _extract_image_urls_from_messages(messages: Any) -> list[str]: - if not isinstance(messages, list): - return [] - image_urls: list[str] = [] - for message in messages: - if not isinstance(message, dict): - continue - content = message.get("content") - if isinstance(content, str): - try: - content = json.loads(content) # Some clients serialize content parts as JSON text. - except json.JSONDecodeError: - continue - if isinstance(content, list): - for content_part in content: - if not isinstance(content_part, dict) or content_part.get("type") != "image_url": - continue - image_url = content_part.get("image_url") - if isinstance(image_url, dict) and isinstance(image_url.get("url"), str): - image_urls.append(image_url["url"]) - return image_urls - - # [multimodal-patch] Recover per-triplet image URLs from raw model_request events. -# Replicates the server triplet view (dedupe by prompt_token_ids, keep last) and the +# Replicates the server triplet view (dedupe by prompt token ids and images, keep last) and the # manager-side filtering below, so the result aligns one-to-one with the kept triplets. # Returns None when the rollout is text-only or when alignment cannot be guaranteed # (caller then leaves image_urls unset). @@ -221,13 +197,14 @@ def _aligned_image_urls(raw_events: list[Event], n_triplets: int) -> list[list[s # Mirror the server-side dedupe: keep the last request for each valid prompt key. # Missing or malformed ids cannot establish that two requests are duplicates. - last_index_by_prompt: dict[tuple[int, ...], int] = {} + last_index_by_prompt: dict[PromptKey, int] = {} kept_indexes: set[int] = set() - for index, (_, prompt_token_ids, _, _) in enumerate(requests): - if not prompt_token_ids or any(type(token_id) is not int for token_id in prompt_token_ids): + for index, (_, prompt_token_ids, _, image_urls) in enumerate(requests): + key = prompt_key(prompt_token_ids, image_urls) + if key is None: kept_indexes.add(index) continue - last_index_by_prompt[tuple(prompt_token_ids)] = index + last_index_by_prompt[key] = index kept_indexes.update(last_index_by_prompt.values()) aligned: list[list[str] | None] = [] diff --git a/docs/80-example-multimodal-qa.md b/docs/80-example-multimodal-qa.md index eb6ecb970..2ea74bc1e 100644 --- a/docs/80-example-multimodal-qa.md +++ b/docs/80-example-multimodal-qa.md @@ -8,6 +8,8 @@ Multimodal QA is a minimal end-to-end example of **multimodal (image) reinforcem The example exercises the multimodal training data path: image URLs are recovered from the raw `model_request` events, aligned with the training triplets, and turned into per-row `multi_modal_inputs` (`pixel_values` + `image_grid_thw`) with mrope `(batch, 4, seq_len)` position ids, so the training forward pass actually consumes the image features. +Within a rollout, requests are deduplicated only when both their prompt token IDs and ordered image URLs match. Different images can produce identical image placeholder tokens, so their transitions are kept separately; repeated requests with the same image URLs retain the last response. + ## Data Preparation The dataset is generated with PIL at startup — there is nothing to download: diff --git a/tests/server/test_endpoints.py b/tests/server/test_endpoints.py index e23fadfae..7df397ccb 100644 --- a/tests/server/test_endpoints.py +++ b/tests/server/test_endpoints.py @@ -314,6 +314,47 @@ def post_model_request(prompt_token_ids: list[int], response_token_ids: list[int assert [event["data"]["response_token_ids"] for event in triplet_events[1:]] == [[20], [30]] +@pytest.mark.parametrize("serialized_content", [False, True]) +def test_triplet_events_distinguish_image_prompts( + client: TestClient, auth_headers: dict[str, str], serialized_content: bool +): + import json + + rollout_id = _rollout(client, auth_headers)["rollout_id"] + image_a = "data:image/png;base64,YQ==" + image_b = "data:image/png;base64,Yg==" + images_by_request = [[image_a], [image_b], [image_a], [image_a, image_b], [image_b, image_a]] + for index, images in enumerate(images_by_request): + content = [{"type": "image_url", "image_url": {"url": url}} for url in images] + response = client.post( + f"/api/rollouts/{rollout_id}/attempt/0/events", + json={ + "event_type": "model_request", + "data": { + "request": { + "messages": [ + {"role": "user", "content": json.dumps(content) if serialized_content else content} + ] + }, + "response": {"prompt_token_ids": [1, 2], "choices": [{"token_ids": [index]}]}, + }, + }, + headers=auth_headers, + ) + assert response.status_code == 200 + + response = client.get(f"/api/rollouts/{rollout_id}/events?format=triplet", headers=auth_headers) + assert response.status_code == 200 + events = response.json() + assert [event["data"]["response_token_ids"] for event in events] == [[1], [2], [3], [4]] + # Image identity is needed only while deduplicating the raw events. + assert all( + set(event["data"]) == {"prompt_token_ids", "response_token_ids", "response_log_probs", "server"} + for event in events + ) + assert "data:image" not in response.text + + @pytest.mark.parametrize("prompt_token_ids", [None, [], [[1]], ["1"], [True], [1.0], "bad-ids", {"token": 1}]) def test_triplet_events_do_not_dedupe_without_valid_prompt_tokens( client: TestClient, diff --git a/tests/verl/test_agl_rollout_manager.py b/tests/verl/test_agl_rollout_manager.py index d1583bb96..c06dd2659 100644 --- a/tests/verl/test_agl_rollout_manager.py +++ b/tests/verl/test_agl_rollout_manager.py @@ -198,6 +198,7 @@ def test_build_completed_rollout_aligns_image_urls_with_triplets() -> None: _event("reward", {"value": 1.0}), ] triplet_events = [ + _triplet_model_request_event([1, 2], [3, 4]), _triplet_model_request_event([1, 2], [30, 40]), _triplet_model_request_event([7, 8], [9]), _triplet_model_request_event([10], [11]), @@ -217,8 +218,8 @@ def test_build_completed_rollout_aligns_image_urls_with_triplets() -> None: ) assert completed.triplets is not None - assert len(completed.triplets) == 3 - assert [triplet.image_urls for triplet in completed.triplets] == [[_IMG2], [_IMG, _IMG2], None] + assert len(completed.triplets) == 4 + assert [triplet.image_urls for triplet in completed.triplets] == [[_IMG], [_IMG2], [_IMG, _IMG2], None] @pytest.mark.parametrize("prompt_token_ids", [None, [], [[1]], ["1"], [True], [1.0], "bad-ids", {"token": 1}]) @@ -267,3 +268,97 @@ def test_aligned_image_urls_text_only_rollout_returns_none_without_warning( # Text-only rollouts keep the exact original behavior: no alignment, no warning. assert _aligned_image_urls(raw_events, 2) is None assert capsys.readouterr().out == "" + + +@pytest.mark.parametrize("routed_experts", [False, True]) +def test_distinct_images_survive_server_to_training_batch(routed_experts: bool) -> None: + import base64 + import io + from typing import cast + + import numpy as np + import torch + from fastapi.testclient import TestClient + from PIL import Image + + from agentlightning.client import AgentLightningSyncClient + from agentlightning.server.app import create_app + from agentlightning.verl.rollout_adapter import RolloutAdapter + + class ImageProcessor: + def __call__(self, text: object, images: list, return_tensors: str) -> dict: + return {"pixel_values": torch.tensor([image.getpixel((0, 0)) for image in images])} + + urls = [] + for color in ("red", "blue"): + buffer = io.BytesIO() + Image.new("RGB", (2, 2), color).save(buffer, format="PNG") + urls.append("data:image/png;base64," + base64.b64encode(buffer.getvalue()).decode()) + + app = create_app( + { + "key": "test", + "default_proxy": {"model_name": "test-model", "train": {"temperature": 1}, "val": {"temperature": 0}}, + } + ) + with TestClient(app, headers={"Authorization": "Bearer test"}) as client: + rollout_id = client.post("/api/rollouts", json=[{"input": {}}]).json()[0]["rollout_id"] + try: + client.patch( + f"/api/rollouts/{rollout_id}", json={"status": {"state": "running", "last_attempt_id": "0"}} + ).raise_for_status() + for index, url in enumerate([urls[0], urls[0], urls[1]]): + event = _raw_model_request_event([url], [1, 2], [index + 3]) + event.data["response"]["choices"][0]["logprobs"] = {"content": [{"logprob": -0.1}]} + if routed_experts: + buffer = io.BytesIO() + np.save(buffer, np.array([index * 10, index * 10 + 1], dtype=np.uint8).reshape(2, 1, 1)) + event.data["routed_experts"] = base64.b64encode(buffer.getvalue()).decode() + response = client.post( + f"/api/rollouts/{rollout_id}/attempt/0/events", + json={"event_type": "model_request", "data": event.data}, + ) + assert response.status_code == 200 + client.post( + f"/api/rollouts/{rollout_id}/attempt/0/events", + json={"event_type": "reward", "data": {"value": 0.75}}, + ).raise_for_status() + client.patch( + f"/api/rollouts/{rollout_id}", json={"status": {"state": "succeeded", "last_attempt_id": "0"}} + ).raise_for_status() + manager = object.__new__(AglRolloutManagerBase) + manager.client = cast(AgentLightningSyncClient, client) + completed = manager._build_completed_rollout( + EnqueuedRollout(data_id="sample", rollout_id=rollout_id, step=0, sample_idx_in_step=0, enqueue_time=0), + manager._get_rollout(rollout_id), + ) + assert completed.triplets is not None + assert [triplet.image_urls for triplet in completed.triplets] == [[urls[0]], [urls[1]]] + + adapter = RolloutAdapter( + max_prompt_length=4, + max_response_length=2, + device=torch.device("cpu"), + pad_token_id=0, + processor=ImageProcessor(), + require_routed_experts=routed_experts, + ) + batch, _ = adapter.get_train_data_batch([completed]) + assert batch.batch is not None + assert batch.batch["responses"].tolist() == [[4, 0], [5, 0]] + assert batch.batch["is_drop_mask"].tolist() == [False, False] + assert batch.batch["token_level_scores"].sum(-1).tolist() == [0.75, 0.75] + torch.testing.assert_close(batch.batch["rollout_log_probs"], torch.tensor([[-0.1, 0.0], [-0.1, 0.0]])) + assert [row["pixel_values"].tolist() for row in batch.non_tensor_batch["multi_modal_inputs"]] == [ + [[255, 0, 0]], + [[0, 0, 255]], + ] + if routed_experts: + assert batch.batch["routed_experts"][:, :, 0, 0].tolist() == [ + [0, 0, 10, 11, 0, 0], + [0, 0, 20, 21, 0, 0], + ] + raw_events = client.get(f"/api/rollouts/{rollout_id}/events").json() + assert ["routed_experts" in event["data"] for event in raw_events[:3]] == [False, True, True] + finally: + client.delete(f"/api/rollouts/{rollout_id}").raise_for_status()