[CI/Bugfix] Async Request ID Aliasing (#3953)

Signed-off-by: Alex Brooks <albrooks@redhat.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

.buildkite/test-merge.yml

@@ -106,43 +106,43 @@ steps:
   - group: ":card_index_dividers: Entrypoints Test"
     depends_on: upload-merge-pipeline
     steps:
-      # - label: "Entrypoint Test with H100"
-      #   timeout_in_minutes: 35
-      #   commands:
-      #     - pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and H100" --run-level "advanced_model"
-      #   agents:
-      #     queue: "mithril-h100-pool"
-      #   plugins:
-      #     - kubernetes:
-      #         podSpec:
-      #           containers:
-      #             - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-      #               resources:
-      #                 limits:
-      #                   nvidia.com/gpu: 2
-      #               volumeMounts:
-      #                 - name: devshm
-      #                   mountPath: /dev/shm
-      #                 - name: hf-cache
-      #                   mountPath: /root/.cache/huggingface
-      #               env:
-      #                 - name: HF_HOME
-      #                   value: /root/.cache/huggingface
-      #                 - name: HF_TOKEN
-      #                   valueFrom:
-      #                     secretKeyRef:
-      #                       name: hf-token-secret
-      #                       key: token
-      #           nodeSelector:
-      #             node.kubernetes.io/instance-type: gpu-h100-sxm
-      #           volumes:
-      #             - name: devshm
-      #               emptyDir:
-      #                 medium: Memory
-      #             - name: hf-cache
-      #               hostPath:
-      #                 path: /mnt/hf-cache
-      #                 type: DirectoryOrCreate
+      - label: "Entrypoint Test with H100"
+        timeout_in_minutes: 35
+        commands:
+          - pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and H100" --run-level "advanced_model"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 2
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                      - name: HF_TOKEN
+                        valueFrom:
+                          secretKeyRef:
+                            name: hf-token-secret
+                            key: token
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
 
       - label: "Entrypoint Test with L4"
         timeout_in_minutes: 30

tests/entrypoints/test_async_omni.py

@@ -1,12 +1,19 @@
 import asyncio
+import re
 from types import SimpleNamespace
 
 import pytest
 from vllm.sampling_params import RequestOutputKind, SamplingParams
 
+from tests.helpers.mark import hardware_test
+from tests.helpers.stage_config import get_deploy_config_path
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+pytestmark = [pytest.mark.core_model]
+
+DIFFUSION_MODEL = "riverclouds/qwen_image_random"
+OMNI_MODEL = "Qwen/Qwen2.5-Omni-7B"
+OMNI_STAGE_CONFIG = get_deploy_config_path("ci/qwen2_5_omni_thinker_only.yaml")
 
 
 async def _noop(**kw):
@@ -59,6 +66,89 @@ def get_async_omni_instance(fake_add_request=_noop, fake_abort_request=_noop) ->
     return omni
 
 
+@pytest.mark.cpu
+def test_generate_submits_randomized_id_to_engine():
+    """Ensure the engine receives a UUID-suffixed ID, not the raw request ID"""
+
+    async def run():
+        submitted_ids = []
+        omni = get_async_omni_instance(fake_add_request=get_fake_add_request(submitted_ids))
+
+        req_id = "my-req-1"
+        async for _ in omni.generate(
+            prompt={"prompt": "test"},
+            request_id=req_id,
+            sampling_params_list=[SimpleNamespace()],
+            output_modalities=["text"],
+        ):
+            pass
+
+        assert len(submitted_ids) == 1
+        assert submitted_ids[0] != req_id
+        assert submitted_ids[0].startswith(f"{req_id}-")
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize(
+    "req_ids,cancel_prefix,expected_cancel_count",
+    [
+        (["cancel-me"], "cancel-me", 1),
+        (["cancel-me", "cancel-me"], "cancel-me", 2),
+        (["cancel-hello", "cancel-hello-world"], "cancel-hello", 1),
+    ],
+)
+def test_abort_handles_internal_request_mapping(req_ids: list[str], cancel_prefix: str, expected_cancel_count: int):
+    """Ensure that abort() with the user-visible ID resolves correctly.
+
+    NOTE: In the case of concurrent / colliding request(s), all requests matching the
+    user provided request ID will be aborted."""
+
+    async def run():
+        aborted_batches = []
+        omni = get_async_omni_instance(
+            fake_abort_request=get_fake_abort(aborted_batches),
+        )
+
+        async def exhaust(agen):
+            async for _ in agen:
+                pass
+
+        tasks = []
+        for user_request_id in req_ids:
+            t = asyncio.create_task(
+                exhaust(
+                    omni.generate(
+                        prompt={"prompt": "test"},
+                        request_id=user_request_id,
+                        sampling_params_list=[SimpleNamespace()],
+                    )
+                )
+            )
+            await asyncio.sleep(0)
+            tasks.append(t)
+
+        assert len(omni.request_states) == len(req_ids)
+        await omni.abort(cancel_prefix)
+
+        assert len(aborted_batches) == 1
+        aborted_ids = aborted_batches[0]
+        # Aborted requests will have fmt {ext_id}-{UUID} to avoid collisions
+        for rid in aborted_ids:
+            assert re.fullmatch(rf"{re.escape(cancel_prefix)}-[0-9a-f]+", rid)
+        assert len(aborted_ids) == expected_cancel_count
+        assert len(set(aborted_ids)) == expected_cancel_count
+        assert len(omni.request_states) == len(req_ids) - expected_cancel_count
+
+        for t in tasks:
+            t.cancel()
+        await asyncio.gather(*tasks, return_exceptions=True)
+
+    asyncio.run(run())
+
+
+@pytest.mark.cpu
 def test_generate_accepts_request_after_repeated_cancellations():
     async def run_test():
         submitted_request_ids = []
@@ -91,22 +181,23 @@ def test_generate_accepts_request_after_repeated_cancellations():
                 await task
 
         assert len(await collect_outputs("after-cancel")) == 1
-        assert submitted_request_ids == [
-            "baseline",
-            "cancel-0",
-            "cancel-1",
-            "cancel-2",
-            "after-cancel",
-        ]
-        assert aborted_request_batches == [
-            ["cancel-0"],
-            ["cancel-1"],
-            ["cancel-2"],
-        ]
+
+        # Check prefixes instead of equality since generate will add
+        # a random numeric suffix to ensure the request ID is unique
+        expected_prefixes = ["baseline-", "cancel-0-", "cancel-1-", "cancel-2-", "after-cancel-"]
+        assert len(submitted_request_ids) == len(expected_prefixes)
+        for submitted, prefix in zip(submitted_request_ids, expected_prefixes):
+            assert submitted.startswith(prefix)
+
+        assert len(aborted_request_batches) == 3
+        for batch, prefix in zip(aborted_request_batches, ["cancel-0-", "cancel-1-", "cancel-2-"]):
+            assert len(batch) == 1
+            assert batch[0].startswith(prefix)
 
     asyncio.run(run_test())
 
 
+@pytest.mark.cpu
 @pytest.mark.parametrize(
     "output_kind", [RequestOutputKind.DELTA, RequestOutputKind.FINAL_ONLY, RequestOutputKind.CUMULATIVE]
 )
@@ -132,3 +223,79 @@ def test_output_kind_is_preserved_with_explicit_sampling_params(output_kind):
 
     asyncio.run(run())
     assert captured_params[0].output_kind == output_kind
+
+
+# End to end tests for ensuring internal manipulation of request ID
+# in diffusion / Omni models don't leak back to the user.
+#
+# NOTE: It seems like we currently need the shutdowns here, otherwise
+# running the tests sequentially seems to leave a zombie process in diffusion
+# that can OOM the Omni tests.
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.omni
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "req_id",
+    ["my-req-1", "img_gen-abc123", "chatcmpl-xyz"],
+    ids=["plain", "prefixed-img", "prefixed-chat"],
+)
+async def test_diffusion_generate_preserves_request_id(req_id):
+    """Ensure diffusion model requests don't leak internal UUID-suffixed req id back to user."""
+    engine = AsyncOmni(model=DIFFUSION_MODEL)
+    try:
+        async for output in engine.generate("a white cat", request_id=req_id):
+            assert output.request_id == req_id
+    finally:
+        engine.shutdown()
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.omni
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "req_id",
+    ["my-req-1", "img_gen-abc123", "chatcmpl-xyz"],
+    ids=["plain", "prefixed-img", "prefixed-chat"],
+)
+async def test_omni_generate_preserves_request_id(req_id):
+    """Ensure omni model requests don't leak internal UUID-suffixed req id back to user."""
+    engine = AsyncOmni(model=OMNI_MODEL, stage_configs_path=OMNI_STAGE_CONFIG)
+    try:
+        async for output in engine.generate(
+            "Say hello in one word.",
+            request_id=req_id,
+            output_modalities=["text"],
+        ):
+            assert output.request_id == req_id
+    finally:
+        engine.shutdown()
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.omni
+@pytest.mark.asyncio
+async def test_diffusion_generate_empty_request_id():
+    """Empty request_id should get a generated internal ID, not stay empty (diffusion)"""
+    engine = AsyncOmni(model=DIFFUSION_MODEL)
+    try:
+        async for output in engine.generate("a white cat", request_id=""):
+            assert output.request_id != ""
+    finally:
+        engine.shutdown()
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.omni
+@pytest.mark.asyncio
+async def test_omni_generate_empty_request_id():
+    """Empty request_id should get a generated internal ID, not stay empty (omni)"""
+    engine = AsyncOmni(model=OMNI_MODEL, stage_configs_path=OMNI_STAGE_CONFIG)
+    try:
+        async for output in engine.generate(
+            "Say hello in one word.",
+            request_id="",
+            output_modalities=["text"],
+        ):
+            assert output.request_id != ""
+    finally:
+        engine.shutdown()

tests/entrypoints/test_omni_entrypoints.py

@@ -16,6 +16,7 @@ from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
 from vllm_omni.engine.async_omni_engine import StageRuntimeInfo
 from vllm_omni.engine.messages import ErrorMessage, OutputMessage
 from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.entrypoints.client_request_state import ClientRequestState
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.entrypoints.omni_base import OmniEngineDeadError
 from vllm_omni.outputs import OmniRequestOutput
@@ -165,6 +166,14 @@ class FakeAsyncOmniEngine:
 def _patch_engine(monkeypatch: pytest.MonkeyPatch, engine: FakeAsyncOmniEngine) -> None:
     monkeypatch.setattr("vllm_omni.entrypoints.omni_base.AsyncOmniEngine", lambda *args, **kwargs: engine)
     monkeypatch.setattr("vllm_omni.entrypoints.omni_base.omni_snapshot_download", lambda model: model)
+    # Don't add random UUIDs to requests calling .generate since we usually
+    # just want to check for present requests anyway, and would need to just
+    # strip the UUID. Explicit checks against the mapping are in tests for
+    # AsyncOmni, or explicitly set the client req states.
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.async_omni.AsyncOmni._get_unique_request_id",
+        staticmethod(lambda request_id: request_id),
+    )
 
 
 def test_direct_omni_with_nullified_parser_only_nulls_untyped_override_fields(
@@ -623,16 +632,23 @@ async def test_async_omni_llm_diffusion_yields_text_stream_then_image(monkeypatc
 async def test_async_omni_abort_forwards_to_engine(monkeypatch: pytest.MonkeyPatch):
     engine = FakeAsyncOmniEngine(stage_metadata=THREE_STAGE_META)
     _patch_engine(monkeypatch, engine)
-
     app = AsyncOmni("dummy-model")
+
+    # Requests internally have a random UUID appended to the
+    # external ID to avoid collisions, so this also tests mapping
+    external_req_id = "req-1"
+    req_id = "req-1-12345678"
     try:
-        app.request_states["req-1"] = object()
-        await app.abort("req-1")
+        app.request_states[req_id] = ClientRequestState(
+            request_id=req_id,
+            external_request_id=external_req_id,
+        )
+        await app.abort(external_req_id)
     finally:
         app.shutdown()
 
-    assert engine.aborted == [["req-1"]]
-    assert "req-1" not in app.request_states
+    assert engine.aborted == [[req_id]]
+    assert external_req_id not in app.request_states
 
 
 @pytest.mark.asyncio

vllm_omni/entrypoints/async_omni.py

@@ -23,6 +23,7 @@ from vllm.pooling_params import PoolingParams
 from vllm.renderers.inputs.preprocess import extract_prompt_components
 from vllm.sampling_params import RequestOutputKind, SamplingParams
 from vllm.tasks import SupportedTask
+from vllm.utils import random_uuid
 from vllm.v1.engine.exceptions import EngineDeadError
 
 from vllm_omni.diffusion.data import OmniACK, OmniSleepTask, OmniWakeTask
@@ -215,6 +216,19 @@ class AsyncOmni(EngineClient, OmniBase):
             return None
         return vllm_config.model_config
 
+    @staticmethod
+    def _get_unique_request_id(external_request_id: str):
+        """Get a random new request ID for this request; at the server level,
+        this is usually set by the calling entrypoint, but in direct calls, we
+        need to set it explicitly since we do not allow empty IDs.
+
+        NOTE: in the upstream vLLM, this is done in the InputProcessor's
+        `assign_request_id`.
+        """
+        uuid = random_uuid()
+        prefix = "" if not external_request_id else f"{external_request_id}-"
+        return f"{prefix}{uuid:.8}"
+
     # ==================== Generate Method ====================
 
     async def generate(
@@ -249,7 +263,8 @@ class AsyncOmni(EngineClient, OmniBase):
         Args:
             prompt: A single prompt **or** a list of prompts.  A list
                 triggers batch mode when the diffusion stage is reached.
-            request_id: Unique identifier for this request.
+            request_id: Unique identifier for this request. If one is not provided,
+                a random one will be generated.
             sampling_params_list: List of SamplingParams, one per stage.
                 Must have the same length as the number of stages.
                 If *None*, uses default sampling params for each stage.
@@ -263,11 +278,17 @@ class AsyncOmni(EngineClient, OmniBase):
         Raises:
             ValueError: If sampling_params_list has incorrect length.
         """
+        # Append a random UUID suffix to the request_id to ensure it is unique
+        # and non-empty, similar to vLLM's input processor. The suffix is used
+        # only for internal tracking throughout the request's life.
+        external_request_id = request_id
+        request_id = self._get_unique_request_id(external_request_id)
+
         # Wait until generation is resumed if the engine is paused
         async with self._pause_cond:
             await self._pause_cond.wait_for(lambda: not self._paused)
 
-        logger.debug(f"[AsyncOmni] generate() called for request {request_id}")
+        logger.debug(f"[AsyncOmni] generate() called for request {external_request_id}")
 
         input_stream_task: asyncio.Task | None = None
         try:
@@ -302,7 +323,11 @@ class AsyncOmni(EngineClient, OmniBase):
                 wall_start_ts,
                 final_stage_id_for_e2e,
             )
-            req_state = ClientRequestState(request_id)
+
+            req_state = ClientRequestState(
+                request_id=request_id,
+                external_request_id=external_request_id,
+            )
             req_state.metrics = metrics
             self.request_states[request_id] = req_state
 
@@ -353,11 +378,11 @@ class AsyncOmni(EngineClient, OmniBase):
         except (asyncio.CancelledError, GeneratorExit):
             if input_stream_task is not None and not input_stream_task.done():
                 input_stream_task.cancel()
-            await self.abort(request_id)
+            await self._abort_internal_requests(request_id)
             logger.info(f"[AsyncOmni] Request {request_id} aborted.")
             raise
         except Exception as e:
-            await self.abort(request_id)
+            await self._abort_internal_requests(request_id)
             logger.info(f"[AsyncOmni] Request {request_id} failed (input error): {e}")
             raise
 
@@ -541,6 +566,8 @@ class AsyncOmni(EngineClient, OmniBase):
             )
 
             if output_to_yield:
+                # Set the external request ID back to the user yielded input
+                output_to_yield.request_id = req_state.external_request_id or output_to_yield.request_id
                 logger.debug(
                     "[AsyncOmni] req=%s stage-%s yielding final_output_type=%s",
                     request_id,
@@ -687,9 +714,26 @@ class AsyncOmni(EngineClient, OmniBase):
     async def abort(self, request_id: str | Iterable[str]) -> None:
         """Abort request(s) via the Orchestrator."""
         request_ids = [request_id] if isinstance(request_id, str) else list(request_id)
+        # Map the external user request IDs to internal IDs used by the Orchestrator.
+        # NOTE: If the user request_id matches multiple requests, all of them will be
+        # aborted. This is also what happens in this case in vLLM's output processor.
+        internal_ids = [s.request_id for s in self.request_states.values() if s.external_request_id in request_ids]
+        await self._abort(internal_ids)
+
+    async def _abort_internal_requests(self, request_id: str | Iterable[str]):
+        """Abort request(s) via the Orchestrator given internal request IDs,
+        which take the format <external_request_id>-<UUID>.
+        """
+        request_ids = [request_id] if isinstance(request_id, str) else list(request_id)
+        # Request IDs are already internal, so we just need to get the matching states.
+        internal_req_ids = [rid for rid in request_ids if rid in self.request_states]
+        await self._abort(internal_req_ids)
+
+    async def _abort(self, request_ids: list[str]) -> None:
+        """Submit request IDs to be aborted to the engine."""
         await self.engine.abort_async(request_ids)
-        for req_id in request_ids:
-            self.request_states.pop(req_id, None)
+        for rid in request_ids:
+            self.request_states.pop(rid, None)
         if self.log_stats:
             logger.info("[AsyncOmni] Aborted request(s) %s", ",".join(request_ids))
 

vllm_omni/entrypoints/client_request_state.py

@@ -6,8 +6,14 @@ from vllm_omni.metrics import OrchestratorAggregator
 class ClientRequestState:
     """Tracks the state of an individual request in the orchestrator."""
 
-    def __init__(self, request_id: str, queue: asyncio.Queue | None = None):
+    def __init__(
+        self,
+        request_id: str,
+        external_request_id: str | None = None,
+        queue: asyncio.Queue | None = None,
+    ):
         self.request_id = request_id
+        self.external_request_id = external_request_id
         self.stage_id: int | None = None
         self.queue = queue if queue is not None else asyncio.Queue()
         self.metrics: OrchestratorAggregator | None = None