Add stage_init_timeout=300 (#1153)

Signed-off-by: princepride <wangzhipeng628@gmail.com>
This commit is contained in:
汪志鹏 authored and GitHub committed 2026-02-02 17:58:48 +08:00
1 parent 9194c1dd84
commit b6c846cc95
3 files changed
+6 -6

No files matched your search

@@ -12,7 +12,7 @@ stage_args:
model_arch: BagelForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.60
gpu_memory_utilization: 0.35
enforce_eager: true
trust_remote_code: true
engine_output_type: text
@@ -44,7 +44,7 @@ stage_args:
max_batch_size: 1
engine_args:
model_stage: dit
gpu_memory_utilization: 0.4
gpu_memory_utilization: 0.55
enforce_eager: true
trust_remote_code: true
engine_output_type: image
@@ -12,7 +12,7 @@ stage_args:
model_arch: BagelForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.65
gpu_memory_utilization: 0.35
enforce_eager: true
trust_remote_code: true
engine_output_type: text
@@ -43,7 +43,7 @@ stage_args:
max_batch_size: 1
engine_args:
model_stage: dit
gpu_memory_utilization: 0.4
gpu_memory_utilization: 0.55
enforce_eager: true
trust_remote_code: true
engine_output_type: image
@@ -161,7 +161,7 @@ def _generate_bagel_image(omni: Omni, prompt: str = DEFAULT_PROMPT) -> Image.Ima
def test_bagel_text2img_shared_memory_connector():
"""Test Bagel text2img with shared memory connector."""
config_path = str(Path(__file__).parent / "stage_configs" / "bagel_sharedmemory_ci.yaml")
omni = Omni(model="ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=config_path)
omni = Omni(model="ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=config_path, stage_init_timeout=300)
try:
generated_image = _generate_bagel_image(omni)
@@ -288,7 +288,7 @@ def test_bagel_text2img_mooncake_connector():
http_port=MOONCAKE_HTTP_PORT,
)
omni = Omni(model="ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=temp_config_file)
omni = Omni(model="ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=temp_config_file, stage_init_timeout=300)
generated_image = _generate_bagel_image(omni)
_validate_pixels(generated_image)