diff --git a/.buildkite/test-merge.yml b/.buildkite/test-merge.yml index 5e75f58ed..318b872b5 100644 --- a/.buildkite/test-merge.yml +++ b/.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 diff --git a/tests/entrypoints/test_async_omni.py b/tests/entrypoints/test_async_omni.py index c8ef5eddf..5d7176490 100644 --- a/tests/entrypoints/test_async_omni.py +++ b/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() diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py index 3612020d4..bbef876d4 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/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 diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 3485f98f7..c07bb2085 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/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 -. + """ + 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)) diff --git a/vllm_omni/entrypoints/client_request_state.py b/vllm_omni/entrypoints/client_request_state.py index 1c9103f79..680ab9838 100644 --- a/vllm_omni/entrypoints/client_request_state.py +++ b/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