From 56ca5d9a8f8779336f6dcdd6f73b0ad020eb77fd Mon Sep 17 00:00:00 2001 From: Chen-Yo Sun Date: Wed, 6 May 2026 15:43:01 -0700 Subject: [PATCH] [BugFix] Forward CLI --tokenizer to per-stage engine configs (#3120) Signed-off-by: Chen-Yo Sun --- .../test_async_omni_engine_stage_init.py | 72 +++++++++++++++++++ vllm_omni/engine/arg_utils.py | 7 +- vllm_omni/engine/async_omni_engine.py | 4 ++ vllm_omni/engine/stage_init_utils.py | 21 ++++++ vllm_omni/entrypoints/cli/serve.py | 1 + 5 files changed, 101 insertions(+), 4 deletions(-) diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 68ead868e..9474f8727 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -563,6 +563,78 @@ def test_initialize_llm_replica_passes_stage_init_timeout_to_complete_stage_hand assert captured_timeout == 302 +def test_async_omni_engine_reads_tokenizer_from_engine_args(monkeypatch): + import vllm_omni.engine.async_omni_engine as engine_mod + + class DummyThread: + def __init__(self, *args, **kwargs): + pass + + def start(self): + pass + + def is_alive(self): + return False + + fake_engine_args = types.SimpleNamespace( + _explicit_fields=frozenset({"tokenizer"}), + explicit_kwargs=lambda: {"tokenizer": "/tokenizer/from-engine-args"}, + ) + + monkeypatch.setattr(engine_mod.threading, "Thread", DummyThread) + monkeypatch.setattr( + AsyncOmniEngine, + "_resolve_stage_configs", + lambda self, model, kwargs: ("dummy-config", [types.SimpleNamespace(engine_args={})]), + ) + monkeypatch.setattr(AsyncOmniEngine, "_wait_for_orchestrator_init", lambda *_, **__: None) + + engine = AsyncOmniEngine("dummy-model", engine_args=fake_engine_args) + + assert engine.tokenizer == "/tokenizer/from-engine-args" + engine.shutdown() + + +def test_build_engine_args_cli_tokenizer_overrides_inferred_base_tokenizer(tmp_path): + from vllm_omni.engine.stage_init_utils import build_engine_args_dict + + stage_cfg = types.SimpleNamespace( + stage_id=0, + stage_type="llm", + engine_args={"model_subdir": "llm"}, + default_sampling_params={}, + ) + + engine_args = build_engine_args_dict( + stage_cfg, + str(tmp_path), + cli_tokenizer="/external/tokenizer", + ) + + assert engine_args["model"] == os.path.join(str(tmp_path), "llm") + assert engine_args["tokenizer"] == "/external/tokenizer" + + +def test_build_engine_args_keeps_stage_owned_tokenizer_subdir(tmp_path): + from vllm_omni.engine.stage_init_utils import build_engine_args_dict + + stage_cfg = types.SimpleNamespace( + stage_id=0, + stage_type="llm", + engine_args={"model_subdir": "llm", "tokenizer_subdir": "tokenizer"}, + default_sampling_params={}, + ) + + engine_args = build_engine_args_dict( + stage_cfg, + str(tmp_path), + cli_tokenizer="/external/tokenizer", + ) + + assert engine_args["model"] == os.path.join(str(tmp_path), "llm") + assert engine_args["tokenizer"] == os.path.join(str(tmp_path), "tokenizer") + + def test_build_stage0_input_processor_uses_omni_input_preprocessor(monkeypatch): import vllm_omni.engine.stage_init_utils as init_mod diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 6c10c7500..1ce4f8d82 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -438,10 +438,9 @@ class OrchestratorArgs: parallel_config: Any = None # === Multi-stage guards === - # --tokenizer is captured here so it does not propagate to every stage - # uniformly (different stages often need different tokenizers, e.g. - # qwen3_omni thinker vs talker). Users wanting a per-stage tokenizer - # should set it in the deploy YAML. + # --tokenizer is captured by the orchestrator and forwarded to stages + # only when the stage does not define tokenizer/tokenizer_subdir itself. + # Users wanting a per-stage tokenizer should set it in the deploy YAML. tokenizer: str | None = None diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 130ed014a..e670a52d8 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -261,6 +261,8 @@ class AsyncOmniEngine: ea_dict.pop("model", None) kwargs = {**ea_dict, **kwargs} + self.tokenizer: str | None = kwargs.get("tokenizer") + # ------------------------------------------------------------------ # # Single-stage mode detection # # ------------------------------------------------------------------ # @@ -485,6 +487,7 @@ class AsyncOmniEngine: stage_cfg, self.model, stage_connector_spec=stage_connector_spec, + cli_tokenizer=getattr(self, "tokenizer", None), ) omni_conn_cfg, omni_from, omni_to = omni_kv_connector if omni_conn_cfg: @@ -654,6 +657,7 @@ class AsyncOmniEngine: stage_cfg, self.model, stage_connector_spec=plan.stage_connector_spec, + cli_tokenizer=getattr(self, "tokenizer", None), ) lock_fds = acquire_device_locks( plan.metadata.stage_id, diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 2f007335b..b6d601de5 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -104,6 +104,18 @@ def _resolve_model_tokenizer_paths(model: str, engine_args: dict[str, Any]) -> s return model +def apply_cli_tokenizer( + engine_args: dict[str, Any], + *, + cli_tokenizer: str | None, + stage_defines_tokenizer: bool, +) -> None: + """Forward CLI tokenizer unless the stage config defines its own.""" + if cli_tokenizer is None or stage_defines_tokenizer: + return + engine_args["tokenizer"] = cli_tokenizer + + def terminate_alive_proc(proc, timeout=5): if proc.is_alive(): proc.terminate() @@ -531,6 +543,7 @@ def build_engine_args_dict( stage_config: Any, model: str, stage_connector_spec: dict[str, Any] | None = None, + cli_tokenizer: str | None = None, ) -> dict[str, Any]: """Build the normalized engine args dict for one stage.""" engine_args = stage_config.engine_args @@ -545,7 +558,15 @@ def build_engine_args_dict( stage_id = stage_config.stage_id engine_args_dict = _to_dict(engine_args) + stage_defines_tokenizer = ( + engine_args_dict.get("tokenizer") is not None or engine_args_dict.get("tokenizer_subdir") is not None + ) model = _resolve_model_tokenizer_paths(model, engine_args_dict) + apply_cli_tokenizer( + engine_args_dict, + cli_tokenizer=cli_tokenizer, + stage_defines_tokenizer=stage_defines_tokenizer, + ) engine_args_dict["model"] = model # Stage id must come from stage config instead of inherited CLI kwargs # (e.g. `--stage-id` defaulting to None). diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 75b21139e..9f9eef2ba 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -628,6 +628,7 @@ def run_headless(args: argparse.Namespace) -> None: stage_cfg, model, stage_connector_spec=stage_connector_spec, + cli_tokenizer=getattr(args, "tokenizer", None), ) # Inject omni KV connector config so the engine runner can initialize the