[BugFix] deploy schema stage runtime overrides for CLI diffusion flags no longer works (#4101)

Signed-off-by: xiaohajiayou <923390377@qq.com>

tests/test_config_factory.py

@@ -1588,6 +1588,197 @@ class TestAuraOmniDeploy:
         assert stages[3].yaml_engine_args["model_arch"] == "Qwen3TTSCode2Wav"
 
 
+class TestDeployCliOverrideFlow:
+    """Test deploy-YAML baselines overridden by CLI runtime overrides."""
+
+    def test_diffusion_deploy_fields_can_be_overridden_by_cli(self, tmp_path):
+        deploy_path = tmp_path / "diffusion_stage.yaml"
+        deploy_path.write_text(
+            """
+async_chunk: false
+stages:
+  - stage_id: 0
+    devices: "4,5"
+    parallel_config:
+      pipeline_parallel_size: 1
+      data_parallel_size: 1
+      tensor_parallel_size: 1
+      enable_expert_parallel: false
+      sequence_parallel_size: 1
+      ulysses_degree: 1
+      ring_degree: 1
+      cfg_parallel_size: 1
+      vae_patch_parallel_size: 1
+      use_hsdp: false
+      hsdp_shard_size: -1
+      hsdp_replicate_size: 1
+    engine_args:
+      cache_backend: cache_dit
+      diffusion_attention_backend: FLASH_ATTN
+      diffusion_kv_cache_dtype: auto
+      step_execution: false
+      vae_use_tiling: false
+      enable_cpu_offload: false
+      max_generated_image_size: 1048576
+      tts_max_instructions_length: 1000
+""",
+            encoding="utf-8",
+        )
+
+        pipeline = PipelineConfig(
+            model_type="test_diffusion",
+            stages=(
+                StagePipelineConfig(
+                    stage_id=0,
+                    model_stage="dit",
+                    execution_type=StageExecutionType.DIFFUSION,
+                    final_output=True,
+                    final_output_type="image",
+                ),
+            ),
+        )
+        deploy = load_deploy_config(deploy_path)
+        stages = merge_pipeline_deploy(pipeline, deploy)
+        stage = stages[0]
+
+        assert stage.yaml_engine_args["parallel_config"]["pipeline_parallel_size"] == 1
+        assert stage.yaml_engine_args["parallel_config"]["data_parallel_size"] == 1
+        assert stage.yaml_engine_args["parallel_config"]["tensor_parallel_size"] == 1
+        assert stage.yaml_engine_args["parallel_config"]["enable_expert_parallel"] is False
+        assert stage.yaml_engine_args["parallel_config"]["sequence_parallel_size"] == 1
+        assert stage.yaml_engine_args["parallel_config"]["ulysses_degree"] == 1
+        assert stage.yaml_engine_args["parallel_config"]["ring_degree"] == 1
+        assert stage.yaml_engine_args["parallel_config"]["cfg_parallel_size"] == 1
+        assert stage.yaml_engine_args["parallel_config"]["vae_patch_parallel_size"] == 1
+        assert stage.yaml_engine_args["cache_backend"] == "cache_dit"
+        assert stage.yaml_engine_args["parallel_config"]["use_hsdp"] is False
+        assert stage.yaml_engine_args["parallel_config"]["hsdp_shard_size"] == -1
+        assert stage.yaml_engine_args["parallel_config"]["hsdp_replicate_size"] == 1
+        assert stage.yaml_engine_args["diffusion_attention_backend"] == "FLASH_ATTN"
+        assert stage.yaml_engine_args["diffusion_kv_cache_dtype"] == "auto"
+        assert stage.yaml_engine_args["step_execution"] is False
+        assert stage.yaml_engine_args["vae_use_tiling"] is False
+        assert stage.yaml_engine_args["enable_cpu_offload"] is False
+        assert stage.yaml_engine_args["max_generated_image_size"] == 1048576
+        assert stage.yaml_engine_args["tts_max_instructions_length"] == 1000
+
+        stage.runtime_overrides = StageConfigFactory._merge_cli_overrides(
+            stage,
+            {
+                "pipeline_parallel_size": 2,
+                "data_parallel_size": 3,
+                "tensor_parallel_size": 4,
+                "enable_expert_parallel": True,
+                "sequence_parallel_size": 24,
+                "ulysses_degree": 2,
+                "ring_degree": 4,
+                "cfg_parallel_size": 2,
+                "vae_patch_parallel_size": 2,
+                "cache_backend": "tea_cache",
+                "use_hsdp": True,
+                "hsdp_shard_size": 8,
+                "hsdp_replicate_size": 2,
+                "diffusion_attention_backend": "SAGE_ATTN",
+                "diffusion_kv_cache_dtype": "fp8",
+                "step_execution": True,
+                "vae_use_tiling": True,
+                "enable_cpu_offload": True,
+                "max_generated_image_size": 2097152,
+                "tts_max_instructions_length": 2000,
+            },
+        )
+
+        omega_config = stage.to_omegaconf()
+
+        assert omega_config.engine_args.cache_backend == "tea_cache"
+        assert omega_config.engine_args.diffusion_attention_backend == "SAGE_ATTN"
+        assert omega_config.engine_args.diffusion_kv_cache_dtype == "fp8"
+        assert omega_config.engine_args.step_execution is True
+        assert omega_config.engine_args.vae_use_tiling is True
+        assert omega_config.engine_args.enable_cpu_offload is True
+        assert omega_config.engine_args.max_generated_image_size == 2097152
+        assert omega_config.engine_args.tts_max_instructions_length == 2000
+        assert omega_config.engine_args.parallel_config.pipeline_parallel_size == 2
+        assert omega_config.engine_args.parallel_config.data_parallel_size == 3
+        assert omega_config.engine_args.parallel_config.tensor_parallel_size == 4
+        assert omega_config.engine_args.parallel_config.enable_expert_parallel is True
+        assert omega_config.engine_args.parallel_config.sequence_parallel_size == 24
+        assert omega_config.engine_args.parallel_config.ulysses_degree == 2
+        assert omega_config.engine_args.parallel_config.ring_degree == 4
+        assert omega_config.engine_args.parallel_config.cfg_parallel_size == 2
+        assert omega_config.engine_args.parallel_config.vae_patch_parallel_size == 2
+        assert omega_config.engine_args.parallel_config.use_hsdp is True
+        assert omega_config.engine_args.parallel_config.hsdp_shard_size == 8
+        assert omega_config.engine_args.parallel_config.hsdp_replicate_size == 2
+
+    def test_llm_deploy_fields_can_be_overridden_by_cli(self, tmp_path):
+        deploy_path = tmp_path / "llm_stage.yaml"
+        deploy_path.write_text(
+            """
+async_chunk: false
+stages:
+  - stage_id: 0
+    devices: "0,1"
+    engine_args:
+      tensor_parallel_size: 1
+      enable_expert_parallel: false
+      gpu_memory_utilization: 0.5
+      max_num_seqs: 16
+      max_num_batched_tokens: 1024
+      max_model_len: 4096
+      enforce_eager: false
+""",
+            encoding="utf-8",
+        )
+
+        pipeline = PipelineConfig(
+            model_type="test_llm",
+            stages=(
+                StagePipelineConfig(
+                    stage_id=0,
+                    model_stage="thinker",
+                    execution_type=StageExecutionType.LLM_AR,
+                    final_output=True,
+                    final_output_type="text",
+                ),
+            ),
+        )
+        deploy = load_deploy_config(deploy_path)
+        stages = merge_pipeline_deploy(pipeline, deploy)
+        stage = stages[0]
+
+        assert stage.yaml_engine_args["tensor_parallel_size"] == 1
+        assert stage.yaml_engine_args["enable_expert_parallel"] is False
+        assert stage.yaml_engine_args["gpu_memory_utilization"] == 0.5
+        assert stage.yaml_engine_args["max_num_seqs"] == 16
+        assert stage.yaml_engine_args["max_num_batched_tokens"] == 1024
+        assert stage.yaml_engine_args["max_model_len"] == 4096
+        assert stage.yaml_engine_args["enforce_eager"] is False
+
+        stage.runtime_overrides = StageConfigFactory._merge_cli_overrides(
+            stage,
+            {
+                "tensor_parallel_size": 2,
+                "enable_expert_parallel": True,
+                "gpu_memory_utilization": 0.9,
+                "max_num_seqs": 32,
+                "max_num_batched_tokens": 2048,
+                "max_model_len": 8192,
+                "enforce_eager": True,
+            },
+        )
+
+        omega_config = stage.to_omegaconf()
+
+        assert omega_config.engine_args.tensor_parallel_size == 2
+        assert omega_config.engine_args.enable_expert_parallel is True
+        assert omega_config.engine_args.gpu_memory_utilization == 0.9
+        assert omega_config.engine_args.max_num_seqs == 32
+        assert omega_config.engine_args.max_num_batched_tokens == 2048
+        assert omega_config.engine_args.max_model_len == 8192
+        assert omega_config.engine_args.enforce_eager is True
+
+
 class TestSentinelDefaultPrecedence:
     """Caller-typed (non-None) values win over YAML; None values fall through
     to YAML / dataclass defaults (#3035)."""

vllm_omni/config/stage_config.py

@@ -63,7 +63,12 @@ def build_stage_runtime_overrides(
     if internal_keys is None:
         from vllm_omni.engine.arg_utils import SHARED_FIELDS, internal_blacklist_keys
 
-        internal_keys = internal_blacklist_keys() | SHARED_FIELDS
+        # Some fields are modeled as orchestrator-owned for top-level CLI
+        # parsing, but are also legitimate deploy-time stage overrides. Keep
+        # the default blacklist for true orchestrator/shared fields while
+        # allowing any field explicitly represented by the deploy schema to
+        # continue flowing into per-stage overrides.
+        internal_keys = (internal_blacklist_keys() | SHARED_FIELDS) - deploy_runtime_override_keys()
 
     result: dict[str, Any] = {}
 
@@ -303,7 +308,7 @@ class StageDeployConfig:
     the top level of ``DeployConfig`` and propagated to every stage.
     """
 
-    # === Omni fields ===
+    # === Omni stage wrapper fields ===
     # Stage identity and Omni runtime placement.
     stage_id: int
     devices: str | None = None
@@ -316,22 +321,33 @@ class StageDeployConfig:
     default_sampling_params: dict[str, Any] | None = None
     subtalker_sampling_params: dict[str, Any] | None = None
 
-    # === vLLM EngineArgs fields ===
-    # Parallelism and scheduler/memory capacity.
+    # === Generic stage engine fields ===
+    # Parallelism, scheduler, and memory-capacity controls.
     tensor_parallel_size: int | None = None
+    enable_expert_parallel: bool | None = None
     gpu_memory_utilization: float | None = None
     max_num_seqs: int | None = None
     max_num_batched_tokens: int | None = None
     max_model_len: int | None = None
 
-    # Execution, scheduling, and KV/cache behavior.
+    # Generic execution, scheduling, and KV/cache behavior.
     enforce_eager: bool | None = None
     async_scheduling: bool | None = None
     disable_hybrid_kv_cache_manager: bool | None = None
     mm_processor_cache_gb: float | None = None
 
-    # Diffusion parallel_config deploy override fields.
-    enable_expert_parallel: bool | None = None
+    # Generic compilation, profiling, tokenizer/config parsing, and model
+    # loading controls.
+    compilation_config: dict[str, Any] | None = None
+    profiler_config: dict[str, Any] | None = None
+    skip_mm_profiling: bool | None = None
+    enable_flashinfer_autotune: bool | None = None
+    config_format: str | None = None
+    load_format: str | None = None
+    tokenizer_mode: str | None = None
+
+    # === Diffusion stage runtime fields ===
+    # Diffusion parallel_config deploy/runtime override fields.
     ulysses_degree: int | None = None
     ulysses_mode: str | None = None
     ring_degree: int | None = None
@@ -343,16 +359,44 @@ class StageDeployConfig:
     hsdp_shard_size: int | None = None
     hsdp_replicate_size: int | None = None
 
-    # Compilation, profiling, tokenizer/config parsing, and model loading.
-    compilation_config: dict[str, Any] | None = None
-    profiler_config: dict[str, Any] | None = None
-    skip_mm_profiling: bool | None = None
-    enable_flashinfer_autotune: bool | None = None
-    config_format: str | None = None
-    load_format: str | None = None
-    tokenizer_mode: str | None = None
+    # Diffusion model loading and adapter construction.
+    model_class_name: str | None = None
+    diffusion_load_format: str | None = None
+    diffusers_load_kwargs: dict[str, Any] | None = None
+    diffusers_call_kwargs: dict[str, Any] | None = None
+    diffusion_quantization_config: str | None = None
+    diffusion_attention_backend: str | None = None
+    diffusion_attention_config: dict[str, Any] | None = None
 
-    # Pass-through vLLM EngineArgs fields that are not represented above.
+    # Diffusion execution, cache, and VAE behavior.
+    cache_backend: str | None = None
+    cache_config: dict[str, Any] | None = None
+    enable_cache_dit_summary: bool | None = None
+    step_execution: bool | None = None
+    vae_use_slicing: bool | None = None
+    vae_use_tiling: bool | None = None
+    boundary_ratio: float | None = None
+    flow_shift: float | None = None
+    diffusion_kv_cache_dtype: str | None = None
+    diffusion_kv_cache_skip_steps: str | None = None
+    diffusion_kv_cache_skip_layers: str | None = None
+    auxiliary_text_encoder: str | None = None
+
+    # Runtime optimizations used by diffusion loading/execution.
+    enable_multithread_weight_load: bool | None = None
+    num_weight_load_threads: int | None = None
+    enable_cpu_offload: bool | None = None
+    enable_layerwise_offload: bool | None = None
+
+    # Diffusion-specific debug and observability knobs.
+    enable_diffusion_pipeline_profiler: bool | None = None
+
+    # Modality/service constraints consumed outside the core engine config.
+    max_generated_image_size: int | None = None
+    tts_max_instructions_length: int | None = None
+
+    # === Pass-through stage engine fields ===
+    # Pass-through stage engine args that are not represented above.
     engine_extras: dict[str, Any] = field(default_factory=dict)
 
 
@@ -409,6 +453,17 @@ _STAGE_RESERVED_KEYS = frozenset(
 _STAGE_DEPLOY_FIELDS = {f.name: f for f in fields(StageDeployConfig) if f.name not in _STAGE_RESERVED_KEYS}
 
 
+def deploy_runtime_override_keys() -> frozenset[str]:
+    """Return deploy-schema fields that are valid CLI/runtime overrides.
+
+    These keys form the positive contract for stage override propagation:
+    stage-scoped deploy knobs plus top-level pipeline-wide engine settings.
+    They must remain overridable even if they are also modeled on
+    ``OrchestratorArgs`` for top-level CLI parsing.
+    """
+    return frozenset(_STAGE_DEPLOY_FIELDS) | frozenset(_PIPELINE_WIDE_ENGINE_FIELDS)
+
+
 def _parse_stage_deploy(stage_data: dict[str, Any]) -> StageDeployConfig:
     """Parse a single stage entry from deploy YAML into StageDeployConfig."""
     # Get the non-reserved keys for this stage