From 3cddda978153f3a19ecd691d343854dd074e587d Mon Sep 17 00:00:00 2001 From: Arifuzzaman Joy <96351147+Arifuzzamanjoy@users.noreply.github.com> Date: Tue, 1 Sep 2026 22:45:28 +0600 Subject: [PATCH] [Bugfix] Prevent build_engine_args_dict from mutating stage_config.engine_args (#6783) Signed-off-by: Arifuzzamanjoy --- .../test_build_engine_args_no_mutate.py | 145 ++++++++++++++++++ vllm_omni/engine/stage_init_utils.py | 2 +- 2 files changed, 146 insertions(+), 1 deletion(-) create mode 100644 tests/engine/test_build_engine_args_no_mutate.py diff --git a/tests/engine/test_build_engine_args_no_mutate.py b/tests/engine/test_build_engine_args_no_mutate.py new file mode 100644 index 000000000..cdc966ed2 --- /dev/null +++ b/tests/engine/test_build_engine_args_no_mutate.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Tests that build_engine_args_dict does not mutate its input.""" + +from __future__ import annotations + +import copy +from types import SimpleNamespace + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_stage_config(engine_args: dict | None = None): + """Build a minimal stage config stub for build_engine_args_dict.""" + if engine_args is None: + engine_args = { + "tensor_parallel_size": None, + "model_subdir": None, + "tokenizer_subdir": None, + } + return SimpleNamespace( + stage_id=0, + stage_type="llm", + engine_args=engine_args, + runtime={}, + final_output=True, + final_output_type=None, + is_comprehension=False, + default_sampling_params={}, + default_pooling_params={}, + engine_input_source=[], + custom_process_input_func=None, + prompt_transform_func=None, + prompt_expand_func=None, + cfg_kv_collect_func=None, + ) + + +@pytest.fixture(autouse=True) +def _patch_platform(monkeypatch): + """Stub out platform checks so tests run on any host.""" + monkeypatch.setattr( + "vllm_omni.engine.stage_init_utils.current_omni_platform", + SimpleNamespace( + is_rocm=lambda: False, + get_omni_ar_worker_cls=lambda: "stub", + get_omni_generation_worker_cls=lambda: "stub", + device_control_env_var="CUDA_VISIBLE_DEVICES", + get_device_count=lambda: 1, + ), + ) + + +class TestBuildEngineArgsDictNoMutate: + """build_engine_args_dict must never modify stage_config.engine_args.""" + + def test_original_dict_unchanged_after_call(self): + engine_args = { + "tensor_parallel_size": None, + "omni_kv_config": {"need_send_cache": True}, + } + original = copy.deepcopy(engine_args) + stage_cfg = _make_stage_config(engine_args) + + from vllm_omni.engine.stage_init_utils import build_engine_args_dict + + build_engine_args_dict(stage_cfg, model="dummy-model") + + assert stage_cfg.engine_args == original, "build_engine_args_dict mutated stage_config.engine_args" + + def test_returned_dict_drops_none_tp(self): + engine_args = {"tensor_parallel_size": None} + stage_cfg = _make_stage_config(engine_args) + + from vllm_omni.engine.stage_init_utils import build_engine_args_dict + + result = build_engine_args_dict(stage_cfg, model="dummy-model") + + assert "tensor_parallel_size" not in result + + def test_returned_dict_keeps_explicit_tp(self): + engine_args = {"tensor_parallel_size": 4} + stage_cfg = _make_stage_config(engine_args) + + from vllm_omni.engine.stage_init_utils import build_engine_args_dict + + result = build_engine_args_dict(stage_cfg, model="dummy-model") + + assert result["tensor_parallel_size"] == 4 + + def test_idempotency(self): + engine_args = { + "tensor_parallel_size": None, + "omni_kv_config": {"stage_id": 0}, + } + stage_cfg = _make_stage_config(engine_args) + + from vllm_omni.engine.stage_init_utils import build_engine_args_dict + + result1 = build_engine_args_dict(stage_cfg, model="dummy-model") + result2 = build_engine_args_dict(stage_cfg, model="dummy-model") + + assert result1 == result2, "Calling build_engine_args_dict twice with the same input produced different results" + + def test_nested_dict_not_shared(self): + """Mutating the returned dict's nested values must not affect the original.""" + engine_args = { + "omni_kv_config": {"need_send_cache": True}, + } + original_kv = copy.deepcopy(engine_args["omni_kv_config"]) + stage_cfg = _make_stage_config(engine_args) + + from vllm_omni.engine.stage_init_utils import build_engine_args_dict + + result = build_engine_args_dict(stage_cfg, model="dummy-model") + + if "omni_kv_config" in result: + result["omni_kv_config"]["injected_key"] = "injected_value" + + assert stage_cfg.engine_args["omni_kv_config"] == original_kv, ( + "Mutating the returned dict's nested value propagated back to stage_config.engine_args" + ) + + def test_inject_omni_kv_does_not_leak_back(self): + """inject_omni_kv_connector_config on the result must not pollute stage_config.""" + engine_args = { + "omni_kv_config": {"need_send_cache": True}, + } + original = copy.deepcopy(engine_args) + stage_cfg = _make_stage_config(engine_args) + + from vllm_omni.engine.stage_init_utils import ( + build_engine_args_dict, + inject_omni_kv_connector_config, + ) + + result = build_engine_args_dict(stage_cfg, model="dummy-model") + connector = ({"host": "localhost", "port": 1234}, "stage-0", "stage-1") + inject_omni_kv_connector_config(result, connector, stage_id=0) + + assert stage_cfg.engine_args == original, ( + "inject_omni_kv_connector_config on the returned dict mutated stage_config.engine_args" + ) diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index b34cdd5da..080337e06 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -1293,7 +1293,7 @@ def build_legacy_engine_args_dict( cli_tokenizer: str | None = None, ) -> dict[str, Any]: """Implement engine-argument building for the legacy stage representation.""" - engine_args_dict = _to_dict(stage_config.engine_args) + engine_args_dict = copy.deepcopy(_to_dict(stage_config.engine_args)) # Legacy configs can materialize an omitted optional TP size as None. # Remove it from the detached adapter dict so the backend default applies # without mutating stage_config.engine_args.