Skip to content

Commit 76a7185

Browse files
committed
[Config Refactor][1/N] Two-Tier Stage Configuration & Independent Stage Launch
Introduce a two-tier configuration system separating internal pipeline topology (Tier-1) from user-configurable runtime params (Tier-2). Tier-1 (Pipeline Topology): - StageConfig dataclass with typed fields for stage identity, DAG connections, worker types, and processing hooks - StageTopology with validate_topology() for integration-time checks - Clean YAML files in stage_topologies/ (no runtime params) Tier-2 (Runtime Config): - StageConfigFactory merges CLI overrides into stage configs - Blocklist approach: all engine-registered CLI args forwarded unless in _INTERNAL_KEYS (no hardcoded allowlist to maintain) - Per-stage overrides via --stage-N-* prefix convention - auto_allocate_memory() for shared-device memory splitting Other changes: - OmegaConf isolated in yaml_util.py (5 wrapper functions) - --stage-id CLI arg for independent stage launch - 34 unit tests across 8 categories - Fix UnboundLocalError in get_final_stage_id_for_e2e - Deprecation docstrings on legacy OmegaConf loading paths Signed-off-by: lishunyang <lishunyang12@163.com>
1 parent a3f2d4c commit 76a7185

File tree

15 files changed

+1651
-52
lines changed

15 files changed

+1651
-52
lines changed

examples/offline_inference/qwen3_omni/end2end.py

Lines changed: 67 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
from vllm.multimodal.image import convert_image_mode
2222
from vllm.utils.argparse_utils import FlexibleArgumentParser
2323

24+
# Import StageConfigFactory for Tier-2 CLI override testing
2425
from vllm_omni.entrypoints.omni import Omni
2526

2627
SEED = 42
@@ -294,11 +295,30 @@ def main(args):
294295
else:
295296
query_result = query_func()
296297

298+
# Build kwargs with Tier-2 CLI overrides
299+
omni_kwargs = {
300+
"stage_configs_path": args.stage_configs_path,
301+
"log_stats": args.log_stats,
302+
"stage_init_timeout": args.stage_init_timeout,
303+
}
304+
305+
# Add Tier-2 CLI overrides if specified
306+
if args.gpu_memory_utilization is not None:
307+
omni_kwargs["gpu_memory_utilization"] = args.gpu_memory_utilization
308+
if args.tensor_parallel_size is not None:
309+
omni_kwargs["tensor_parallel_size"] = args.tensor_parallel_size
310+
if args.devices is not None:
311+
omni_kwargs["devices"] = args.devices
312+
if args.enforce_eager:
313+
omni_kwargs["enforce_eager"] = args.enforce_eager
314+
if args.trust_remote_code:
315+
omni_kwargs["trust_remote_code"] = args.trust_remote_code
316+
if args.stage_id is not None:
317+
omni_kwargs["stage_id"] = args.stage_id
318+
297319
omni_llm = Omni(
298320
model=model_name,
299-
stage_configs_path=args.stage_configs_path,
300-
log_stats=args.log_stats,
301-
stage_init_timeout=args.stage_init_timeout,
321+
**omni_kwargs,
302322
)
303323

304324
thinker_sampling_params = SamplingParams(
@@ -458,6 +478,12 @@ def parse_args():
458478
default="output_audio",
459479
help="[Deprecated] Output wav directory (use --output-dir).",
460480
)
481+
parser.add_argument(
482+
"--output-dir",
483+
type=str,
484+
default=None,
485+
help="Output directory for generated files (text and audio).",
486+
)
461487
parser.add_argument(
462488
"--num-prompts",
463489
type=int,
@@ -474,7 +500,44 @@ def parse_args():
474500
"--stage-configs-path",
475501
type=str,
476502
default=None,
477-
help="Path to a stage configs file.",
503+
help="Path to a stage configs file. If not specified, uses auto-detected Tier-1 topology.",
504+
)
505+
# Tier-2 CLI override arguments
506+
parser.add_argument(
507+
"--gpu-memory-utilization",
508+
type=float,
509+
default=None,
510+
help="GPU memory utilization for all stages (Tier-2 override). Example: 0.9",
511+
)
512+
parser.add_argument(
513+
"--tensor-parallel-size",
514+
type=int,
515+
default=None,
516+
help="Tensor parallel size for all stages (Tier-2 override). Example: 2",
517+
)
518+
parser.add_argument(
519+
"--devices",
520+
type=str,
521+
default=None,
522+
help="Device assignment for stages (Tier-2 override). Example: '0,1'",
523+
)
524+
parser.add_argument(
525+
"--enforce-eager",
526+
action="store_true",
527+
default=False,
528+
help="Enforce eager mode for all stages (Tier-2 override).",
529+
)
530+
parser.add_argument(
531+
"--trust-remote-code",
532+
action="store_true",
533+
default=False,
534+
help="Trust remote code for model loading (Tier-2 override).",
535+
)
536+
parser.add_argument(
537+
"--stage-id",
538+
type=int,
539+
default=None,
540+
help="Launch only the specified stage ID for independent stage testing.",
478541
)
479542
parser.add_argument(
480543
"--video-path",

0 commit comments

Comments
 (0)