[Bugfix][EPD] Preserve raw fallback identity and transfer lifecycle

examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py

@@ -29,6 +29,7 @@
 import time
 import uuid
 from collections.abc import AsyncIterator
+from copy import deepcopy
 from typing import Any
 
 import aiohttp
@@ -136,7 +137,6 @@ def rewrite_for_decode(req_data: dict, item_meta: dict[int, dict]) -> dict:
     here -- a second derivation could disagree with the encoder's.
     """
     rewritten = 0
-    transfer_items = []
     idx = 0
     new_messages = []
     for msg in req_data.get("messages", []):
@@ -152,8 +152,6 @@ def rewrite_for_decode(req_data: dict, item_meta: dict[int, dict]) -> dict:
             meta = dict(item_meta.get(idx) or {})
             idx += 1
             item_uuid = meta.pop("mm_hash", None)
-            ec_mm_hash = meta.pop("ec_mm_hash", None) or item_uuid
-            transfer_id = meta.pop("transfer_id", None)
             # Whatever keys the encoder reported are the metadata its model
             # declared as needed to size the placeholder range; the proxy does
             # not need to know their names.
@@ -167,35 +165,19 @@ def rewrite_for_decode(req_data: dict, item_meta: dict[int, dict]) -> dict:
                 for k, v in meta.items()
             }
             if not metadata or not item_uuid:
-                # Nothing to size the placeholder range with. A processor cache
-                # hit is not a cause on its own: with the default `lru` type the
-                # engine restores the item before the scheduler reports it. It
-                # goes missing when the encode request failed, or under
-                # `--mm-processor-cache-type shm`, where a hit replaces the item
-                # with its shared-memory address and only the worker restores
-                # it. Send the media so the decoder can derive the grid itself.
-                new_content.append(item)
-                continue
+                # The parser cannot mix raw and embeds for one modality.
+                return req_data
             embeds_type = EMBEDS_TYPES[item["type"]]
             new_content.append(
                 {"type": embeds_type, embeds_type: metadata, "uuid": item_uuid}
             )
-            if transfer_id is not None:
-                transfer_items.append(
-                    {"mm_hash": ec_mm_hash, "transfer_id": transfer_id}
-                )
             rewritten += 1
         new_messages.append({**msg, "content": new_content})
 
     if not rewritten:
         return req_data
     logger.info("Rewrote %d media item(s) as metadata references", rewritten)
-    rewritten_request = {**req_data, "messages": new_messages}
-    if transfer_items:
-        ec_transfer_params = dict(req_data.get("ec_transfer_params") or {})
-        ec_transfer_params["ec_items"] = transfer_items
-        rewritten_request["ec_transfer_params"] = ec_transfer_params
-    return rewritten_request
+    return {**req_data, "messages": new_messages}
 
 
 def extract_mm_items(request_data: dict) -> list[dict]:
@@ -246,7 +228,6 @@ async def fanout_encoder_primer(
 
     tasks = []
     item_uuids: dict[int, str] = {}
-    item_transfer_ids: dict[int, str] = {}
     item_meta: dict[int, dict] = {}
     ec_params: dict[str, Any] = {}
 
@@ -265,14 +246,14 @@ async def fanout_encoder_primer(
         headers = {"x-request-id": child_req_id, "Content-Type": "application/json"}
 
         # With --no-rewrite the decoder still receives the raw image and derives
-        # the cache key by hashing it, so the encoder must do the same -- passing
-        # a uuid here would make the two disagree and silently defeat the EC
-        # transfer, leaving the decoder to encode the image itself.
-        item_uuid = None if NO_REWRITE else content_uuid(item)
+        # its identity from the original UUID or content. Do not inject a new
+        # UUID on only the encoder side.
+        item_uuid = item.get("uuid")
+        if item_uuid is None and not NO_REWRITE:
+            item_uuid = content_uuid(item)
         if item_uuid is not None:
             item_uuids[idx] = item_uuid
         transfer_id = uuid.uuid4().hex
-        item_transfer_ids[idx] = transfer_id
 
         encoder_req = {
             "model": orig_request.get("model"),
@@ -304,7 +285,8 @@ async def fanout_encoder_primer(
             # would fail the connector's hash match and make the producer
             # invent its own transfer id, which the consumer could then never
             # cancel. Omitting it lets the connector match by position, which
-            # is exact: one encoder request carries exactly one item.
+            # names the first feature. The encoder reports any additional
+            # bindings if this source item expands into multiple features.
             encoder_req["ec_transfer_params"] = {
                 "consumer_zmq": consumer_zmq,
                 "ec_items": [{"transfer_id": transfer_id}],
@@ -349,44 +331,30 @@ async def fanout_encoder_primer(
                 detail=f"Encoder request failed: {detail}",
             )
 
-        # The encoder reports each mm_hash's metadata (e.g. the grid) here,
-        # keyed by the same uuid this proxy assigned above.
+        # Metadata is keyed by the encoder's feature identifiers, which may
+        # be derived from source UUIDs and processing options.
         try:
             params = msgspec.json.decode(await r.read()).get("ec_transfer_params") or {}
         except Exception:
             logger.warning("[%s] Could not read encoder metadata #%d", req_id, idx)
             params = {}
-        if params:
-            # One encoder request carries exactly one item, so there is a
-            # single reported entry. Do not key it by this proxy's uuid: when
-            # media_io_kwargs or mm_processor_kwargs are set the engine
-            # re-hashes the uuid together with them, so the encoder's own
-            # `mm_features[i].identifier` is a derived value this proxy cannot
-            # predict. Fall back to the sole entry, and carry the key the
-            # encoder actually used through as `ec_mm_hash`.
-            ec_mm_hash = item_uuids.get(idx)
-            reported = params.get(ec_mm_hash)
-            if reported is None and len(params) == 1:
-                ((ec_mm_hash, reported),) = params.items()
-            if reported:
-                metadata = reported.get("metadata") or {}
-                if metadata:
-                    item_meta[idx] = {
-                        **metadata,
-                        "mm_hash": item_uuids.get(idx, ec_mm_hash),
-                        "ec_mm_hash": ec_mm_hash,
-                    }
-                    if idx in item_transfer_ids:
-                        item_meta[idx]["transfer_id"] = item_transfer_ids[idx]
-                # Whatever the encoder reported alongside `metadata` is the
-                # connector's own handle on the published embedding (for NIXL,
-                # peer_host/peer_port/size_bytes). The decoder's connector
-                # looks it up by mm_hash on the request, so carry it through.
-                ec_params[ec_mm_hash] = reported
-                if NO_REWRITE and consumer_zmq is not None:
-                    ec_params.setdefault("ec_items", []).append(
-                        {"mm_hash": ec_mm_hash, "transfer_id": item_transfer_ids[idx]}
-                    )
+        # Connector handles remain useful even when placeholder metadata is
+        # unavailable or one source item expands into several features.
+        bindings = params.pop("ec_items", [])
+        ec_params.update(params)
+        if bindings:
+            ec_params.setdefault("ec_items", []).extend(bindings)
+        # The child request identifies the source even when its hash is
+        # derived. Multi-feature expansion needs a combined placeholder
+        # contract; selecting just one entry would drop the other features.
+        if len(params) == 1:
+            ((ec_mm_hash, reported),) = params.items()
+            metadata = reported.get("metadata") or {}
+            if metadata:
+                item_meta[idx] = {
+                    **metadata,
+                    "mm_hash": item_uuids.get(idx, ec_mm_hash),
+                }
 
     logger.info(
         "[%s] All %d encoder requests completed successfully", req_id, len(mm_items)
@@ -578,6 +546,13 @@ async def prepare_for_decode(
     `req_data` is left untouched so a retry starts from the original media
     rather than from a body whose images are already metadata references.
     """
+    if not NO_REWRITE:
+        # Share source identities with E even if the response cannot be
+        # rewritten. Never feed a processed feature hash back as a raw UUID.
+        req_data = deepcopy(req_data)
+        for item in extract_mm_items(req_data):
+            if item.get("uuid") is None:
+                item["uuid"] = content_uuid(item)
     _t0 = time.perf_counter()
     item_meta, ec_params = await fanout_encoder_primer(
         req_data, e_urls, req_id, consumer_zmq

tests/v1/ec_connector/unit/test_epd_proxy_retry.py

@@ -130,7 +130,16 @@ def test_raw_media_keeps_encoder_transfer_identity(proxy, monkeypatch):
     handle = {"metadata": {"image_grid_thw": [1, 2, 2]}}
     monkeypatch.setattr(proxy, "NO_REWRITE", True)
     monkeypatch.setattr(
-        proxy, "encode_session", _EncoderSession([{"encoded-hash": handle}])
+        proxy,
+        "encode_session",
+        _EncoderSession(
+            [
+                {
+                    "encoded-hash": handle,
+                    "ec_items": [{"mm_hash": "encoded-hash", "transfer_id": "actual"}],
+                }
+            ]
+        ),
     )
     body = {
         "messages": [
@@ -324,3 +333,152 @@ def test_pooled_connections_are_retired_before_the_server_closes_them(
         assert proxy.decode_session.connector._keepalive_timeout == pooled_for
     finally:
         asyncio.run(proxy.on_shutdown())
+
+
+@pytest.mark.parametrize("source_uuid", [None, "user-media"])
+@pytest.mark.parametrize("expanded", [False, True])
+def test_fallback_preserves_source_identity_and_all_handles(
+    proxy, monkeypatch, source_uuid, expanded
+):
+    """Raw PD preprocessing must start with the same inputs as E."""
+    from copy import deepcopy
+
+    item = {"type": "video_url", "video_url": {"url": "video-with-audio"}}
+    if source_uuid is not None:
+        item["uuid"] = source_uuid
+    body = {
+        "messages": [{"role": "user", "content": [item]}],
+        "mm_processor_kwargs": {"use_audio_in_video": expanded},
+        "media_io_kwargs": {"video": {"num_frames": 8}},
+    }
+    original = deepcopy(body)
+    hashes = ["video-feature", "audio-feature"] if expanded else ["video-feature"]
+    bindings = [{"mm_hash": h, "transfer_id": f"transfer-{h}"} for h in hashes]
+    handles: dict = {h: {"metadata": {}} for h in hashes}
+    if expanded:
+        handles["video-feature"]["metadata"] = {"video_grid_thw": [1, 2, 2]}
+    seen = []
+
+    async def post(url, data, headers):
+        seen.append(msgspec.json.decode(data))
+        return _EncoderResponse({**handles, "ec_items": bindings})
+
+    session = _EncoderSession([])
+    monkeypatch.setattr(session, "post", post)
+    monkeypatch.setattr(proxy, "encode_session", session)
+    prepared, _, _ = asyncio.run(
+        proxy.prepare_for_decode(body, "r", ["http://e"], "", "tcp://pd")
+    )
+    encoded = seen[0]["messages"][0]["content"][0]
+    forwarded = prepared["messages"][0]["content"][0]
+    assert encoded["uuid"] == (source_uuid or proxy.content_uuid(item))
+    assert forwarded == encoded
+    assert prepared["mm_processor_kwargs"] == seen[0]["mm_processor_kwargs"]
+    assert prepared["media_io_kwargs"] == seen[0]["media_io_kwargs"]
+    assert prepared["ec_transfer_params"] == {**handles, "ec_items": bindings}
+    assert body == original
+
+
+def test_missing_metadata_falls_back_for_the_whole_request(proxy):
+    body = {
+        "messages": [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": "first"},
+                        "uuid": "first",
+                    },
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": "second"},
+                        "uuid": "second",
+                    },
+                ],
+            }
+        ]
+    }
+    assert (
+        proxy.rewrite_for_decode(
+            body,
+            {
+                0: {"mm_hash": "first", "image_grid_thw": [1, 2, 2]},
+            },
+        )
+        == body
+    )
+
+
+def test_rewrite_matches_derived_feature_hash_with_transfer_bindings(
+    proxy, monkeypatch
+):
+    """Transfer entries are not features; child boundaries identify the image."""
+    source_uuid = "user-image"
+    binding = {"mm_hash": "processed-image", "transfer_id": "actual-transfer"}
+    monkeypatch.setattr(
+        proxy,
+        "encode_session",
+        _EncoderSession(
+            [
+                {
+                    "processed-image": {"metadata": {"image_grid_thw": [[1, 2, 2]]}},
+                    "ec_items": [binding],
+                }
+            ]
+        ),
+    )
+    body = {
+        "messages": [
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": "image"},
+                        "uuid": source_uuid,
+                    }
+                ],
+            }
+        ],
+        "mm_processor_kwargs": {"max_pixels": 262144},
+    }
+    prepared, _, _ = asyncio.run(
+        proxy.prepare_for_decode(body, "r", ["http://e"], "", "tcp://pd")
+    )
+    assert prepared["messages"][0]["content"] == [
+        {
+            "type": "image_embeds",
+            "image_embeds": {"image_grid_thw": [1, 2, 2]},
+            "uuid": source_uuid,
+        }
+    ]
+    assert prepared["ec_transfer_params"]["ec_items"] == [binding]
+
+
+def test_repeated_images_keep_distinct_transfers_and_both_rewrite(proxy, monkeypatch):
+    """A hash-keyed handle alone loses one of two actual encoder transfers."""
+    bindings = [{"mm_hash": "same-feature", "transfer_id": t} for t in ("a", "b")]
+    monkeypatch.setattr(
+        proxy,
+        "encode_session",
+        _EncoderSession(
+            [
+                {
+                    "same-feature": {"metadata": {"image_grid_thw": [1, 2, 2]}},
+                    "ec_items": [binding],
+                }
+                for binding in bindings
+            ]
+        ),
+    )
+    item = {"type": "image_url", "image_url": {"url": "same-image"}}
+    body = {"messages": [{"role": "user", "content": [item, item]}]}
+    prepared, _, _ = asyncio.run(
+        proxy.prepare_for_decode(body, "r", ["http://e"], "", "tcp://pd")
+    )
+    assert [i["type"] for i in prepared["messages"][0]["content"]] == [
+        "image_embeds",
+        "image_embeds",
+    ]
+    assert prepared["ec_transfer_params"]["ec_items"] == bindings

tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py

@@ -101,7 +101,7 @@ def test_mooncake_independent_prefill_routing_fails_fast(proxy):
 
 
 def test_decode_rewrite_preserves_engine_reported_ec_hash(proxy):
-    request = {
+    request: dict = {
         "messages": [
             {
                 "role": "user",
@@ -110,13 +110,14 @@ def test_decode_rewrite_preserves_engine_reported_ec_hash(proxy):
         ]
     }
 
+    request["ec_transfer_params"] = {
+        "ec_items": [{"mm_hash": "engine-derived-hash", "transfer_id": "transfer"}]
+    }
     rewritten = proxy.rewrite_for_decode(
         request,
         {
             0: {
                 "mm_hash": "proxy-uuid",
-                "ec_mm_hash": "engine-derived-hash",
-                "transfer_id": "transfer",
                 "image_grid_thw": [1, 2, 3],
             }
         },