[Bugfix] Route text-only chat as per-request comprehension in HunyuanImage3 AR sampler (#6111)

Signed-off-by: MrlixiangWE <mrdanaer@gmail.com>
Co-authored-by: zijianc2 <157244773+zijianc2@users.noreply.github.com>

tests/worker/test_model_sampler_extra_args.py

@@ -0,0 +1,91 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Unit tests for per-request model sampler arguments.
+
+Custom model samplers that set ``model_sampler_wants_extra_args = True``
+(e.g. HunyuanImage3's per-request ``ar_task_mode``, #6088) receive one
+``SamplingParams.extra_args`` entry per batch row, in ``input_batch.req_ids``
+order — the same order ``_build_model_sampler_output_token_ids`` uses. Pure CPU.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.worker.sampling_utils import build_model_sampler_extra_args, call_model_sampler
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _state(extra_args):
+    return SimpleNamespace(sampling_params=SimpleNamespace(extra_args=extra_args))
+
+
+def test_row_order_follows_req_ids():
+    input_batch = SimpleNamespace(req_ids=["b", "a"])
+    requests = {
+        "a": _state({"ar_task_mode": "comprehension"}),
+        "b": _state(None),
+    }
+    assert build_model_sampler_extra_args(input_batch, requests) == [None, {"ar_task_mode": "comprehension"}]
+
+
+def test_missing_request_state_yields_none():
+    assert build_model_sampler_extra_args(SimpleNamespace(req_ids=["gone"]), {}) == [None]
+
+
+def test_request_without_sampling_params_yields_none():
+    input_batch = SimpleNamespace(req_ids=["r0"])
+    assert build_model_sampler_extra_args(input_batch, {"r0": SimpleNamespace(sampling_params=None)}) == [None]
+
+
+def test_empty_batch():
+    assert build_model_sampler_extra_args(SimpleNamespace(req_ids=[]), {}) == []
+
+
+def test_opted_in_sampler_receives_per_request_extra_args(mocker):
+    model = SimpleNamespace(model_sampler_wants_extra_args=True)
+    model_sample = mocker.Mock(return_value="sampled")
+    input_batch = SimpleNamespace(req_ids=["a", "b"])
+    logits = torch.empty(0)
+    metadata = SimpleNamespace()
+    requests = {
+        "a": _state(None),
+        "b": _state({"ar_task_mode": "comprehension"}),
+    }
+
+    result = call_model_sampler(
+        model,
+        model_sample,
+        logits,
+        metadata,
+        input_batch=input_batch,
+        requests=requests,
+    )
+
+    assert result == "sampled"
+    model_sample.assert_called_once_with(
+        logits,
+        metadata,
+        per_req_extra_args=[None, {"ar_task_mode": "comprehension"}],
+    )
+
+
+def test_non_opted_in_sampler_keeps_two_argument_contract(mocker):
+    model = SimpleNamespace(model_sampler_wants_extra_args=False)
+    model_sample = mocker.Mock(return_value="sampled")
+    logits = torch.empty(0)
+    metadata = SimpleNamespace()
+
+    result = call_model_sampler(
+        model,
+        model_sample,
+        logits,
+        metadata,
+        input_batch=SimpleNamespace(req_ids=["a"]),
+        requests={"a": _state({"ar_task_mode": "comprehension"})},
+    )
+
+    assert result == "sampled"
+    model_sample.assert_called_once_with(logits, metadata)