anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Config] Remove diffusion_batch_size and use max_num_seqs instead. (#6484)

Browse Source 
Signed-off-by: LyxWxj <yuxuanlou621@foxmail.com>
Signed-off-by: Yuxuan Rho <104440472+LyxWxj@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
Yuxuan RhoandHongsheng Liu authored and GitHub committed 2026-09-01 01:02:54 +00:00
1 parent 8e26dce1f4
commit 7992c2213e
16 files changed
+36 -95
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/engine/test_async_omni_engine_stage_init.py
+7 -29
View File
@@ -172,7 +172,6 @@ def _make_stage_runtime() -> StageRuntime:
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=1,
diffusion_batch_size=1,
async_chunk=False,
)
@@ -196,7 +195,6 @@ def test_async_omni_engine_initialize_stages_passes_log_stats_to_runtime(monkeyp
engine.model = "dummy-model"
engine.config_path = "dummy-config"
engine.single_stage_mode = False
engine.diffusion_batch_size = 1
engine.async_chunk = False
engine.tokenizer = None
engine._single_stage_id_filter = None
@@ -282,7 +280,6 @@ def test_build_logical_stage_init_plans_handles_stage_without_devices(monkeypatc
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=1,
diffusion_batch_size=1,
async_chunk=False,
)
@@ -371,7 +368,6 @@ def test_initialize_local_diffusion_replica_restores_device_visibility_after_loc
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=1,
diffusion_batch_size=1,
async_chunk=False,
)
@@ -416,7 +412,6 @@ def test_initialize_local_diffusion_replica_passes_stage_init_timeout_and_inline
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=1,
diffusion_batch_size=4,
async_chunk=False,
)
@@ -429,7 +424,6 @@ def test_initialize_local_diffusion_replica_passes_stage_init_timeout_and_inline
def _capture_launch_diffusion_stage_replica(**kwargs):
captured["stage_id"] = kwargs["metadata"].stage_id
captured["stage_init_timeout"] = kwargs["stage_init_timeout"]
captured["batch_size"] = kwargs["batch_size"]
captured["use_inline"] = kwargs["use_inline"]
captured["omni_master_server"] = kwargs["omni_master_server"]
return types.SimpleNamespace(), StageReplicaResources()
@@ -441,7 +435,6 @@ def test_initialize_local_diffusion_replica_passes_stage_init_timeout_and_inline
assert captured == {
"stage_id": 0,
"stage_init_timeout": 302,
"batch_size": 4,
"use_inline": expected_inline,
"omni_master_server": None,
}
@@ -511,21 +504,20 @@ def test_initialize_local_llm_replica_scopes_runtime_env(monkeypatch):
assert os.environ[runtime_env_var] == "parent-value"
def test_initialize_diffusion_stage_forwards_batch_size_without_writing_max_num_seqs(monkeypatch):
def test_initialize_diffusion_stage_preserves_configured_max_num_seqs(monkeypatch):
import vllm_omni.diffusion.stage_diffusion_client as client_mod
import vllm_omni.engine.stage_init_utils as init_mod
od_config = types.SimpleNamespace(max_num_seqs=1)
od_config = types.SimpleNamespace(max_num_seqs=4)
captured: dict[str, object] = {}
monkeypatch.setattr(init_mod, "build_diffusion_config", lambda *args: od_config)
def _capture_client(model, config, metadata, stage_init_timeout, batch_size, use_inline):
def _capture_client(model, config, metadata, stage_init_timeout, use_inline):
captured.update(
model=model,
config=config,
metadata=metadata,
stage_init_timeout=stage_init_timeout,
batch_size=batch_size,
use_inline=use_inline,
)
return object()
@@ -539,27 +531,25 @@ def test_initialize_diffusion_stage_forwards_batch_size_without_writing_max_num_
types.SimpleNamespace(),
metadata,
stage_init_timeout=12,
batch_size=4,
use_inline=True,
)
assert od_config.max_num_seqs == 1
assert od_config.max_num_seqs == 4
assert captured == {
"model": "dummy-model",
"config": od_config,
"metadata": metadata,
"stage_init_timeout": 12,
"batch_size": 4,
"use_inline": True,
}
def test_launch_diffusion_stage_replica_does_not_write_max_num_seqs_from_batch_size(monkeypatch):
def test_launch_diffusion_stage_replica_preserves_configured_max_num_seqs(monkeypatch):
import vllm_omni.diffusion.stage_diffusion_client as client_mod
import vllm_omni.diffusion.stage_diffusion_proc as proc_mod
import vllm_omni.engine.stage_engine_startup as startup_mod
od_config = types.SimpleNamespace(max_num_seqs=1, parallel_config=types.SimpleNamespace(world_size=1))
od_config = types.SimpleNamespace(max_num_seqs=4, parallel_config=types.SimpleNamespace(world_size=1))
monkeypatch.setattr(startup_mod, "build_diffusion_config", lambda *args: od_config)
monkeypatch.setattr(startup_mod, "acquire_device_locks", lambda *args: [])
monkeypatch.setattr(
@@ -596,13 +586,12 @@ def test_launch_diffusion_stage_replica_does_not_write_max_num_seqs_from_batch_s
stage_config=types.SimpleNamespace(),
metadata=types.SimpleNamespace(stage_id=0),
stage_init_timeout=12,
batch_size=4,
use_inline=False,
omni_master_server=omni_master_server,
)
assert result is sentinel_client
assert od_config.max_num_seqs == 1
assert od_config.max_num_seqs == 4
assert resources.manager is proc_manager
@@ -621,7 +610,6 @@ def test_initialize_diffusion_stage_does_not_write_max_num_seqs(monkeypatch):
types.SimpleNamespace(),
metadata,
stage_init_timeout=12,
batch_size=4,
use_inline=True,
)
@@ -672,7 +660,6 @@ def test_launch_diffusion_stage_replica_preserves_step_execution_max_num_seqs(mo
stage_config=types.SimpleNamespace(),
metadata=types.SimpleNamespace(stage_id=0),
stage_init_timeout=12,
batch_size=4,
use_inline=False,
omni_master_server=omni_master_server,
)
@@ -688,7 +675,6 @@ def test_stage_runtime_initializes_stage_pools(monkeypatch):
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=1,
diffusion_batch_size=1,
async_chunk=False,
)
@@ -754,7 +740,6 @@ def test_stage_runtime_passes_log_stats_to_llm_replica_launch(monkeypatch):
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=1,
diffusion_batch_size=1,
async_chunk=False,
log_stats=True,
)
@@ -798,7 +783,6 @@ def test_stage_runtime_passes_log_stats_to_output_processor(monkeypatch):
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=1,
diffusion_batch_size=1,
async_chunk=False,
log_stats=True,
)
@@ -851,7 +835,6 @@ def test_build_logical_stage_init_plans_applies_replica_device_splits(monkeypatc
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=1,
diffusion_batch_size=1,
async_chunk=False,
)
@@ -892,7 +875,6 @@ def test_initialize_stage_replicas_collects_results_by_stage_and_replica_id(monk
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=123,
diffusion_batch_size=1,
async_chunk=False,
)
@@ -930,7 +912,6 @@ def test_remote_replicas_use_distinct_init_group_keys():
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=123,
diffusion_batch_size=1,
async_chunk=False,
)
plan = _make_llm_plan(
@@ -956,7 +937,6 @@ def test_initialize_stages_cleans_up_successful_replicas_after_partial_multi_rep
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=1,
diffusion_batch_size=1,
async_chunk=False,
)
@@ -993,7 +973,6 @@ def test_initialize_stages_cleans_up_late_successful_replicas_after_early_multi_
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=1,
diffusion_batch_size=1,
async_chunk=False,
)
@@ -1032,7 +1011,6 @@ def test_initialize_local_llm_replica_passes_stage_init_timeout_to_complete_stag
model="dummy-model",
config_path="dummy-config",
stage_init_timeout=1,
diffusion_batch_size=1,
async_chunk=False,
)
stage_init_timeout = 302
tests/engine/test_single_stage_mode.py
+2 -13
View File
@@ -575,7 +575,6 @@ class TestSingleStageInitialization:
model="fake-model",
config_path="/fake/stages.yaml",
stage_init_timeout=60,
diffusion_batch_size=2,
async_chunk=False,
single_stage_id_filter=stage_id_filter,
omni_master_address="127.0.0.1",
@@ -807,7 +806,6 @@ class TestSingleStageInitialization:
model="fake-model",
config_path="/fake/stages.yaml",
stage_init_timeout=60,
diffusion_batch_size=2,
async_chunk=False,
)
mocker.patch.object(runtime_mod, "prepare_engine_environment")
@@ -860,7 +858,6 @@ class TestSingleStageReplicaInitialization:
model="fake-model",
config_path="/fake/stages.yaml",
stage_init_timeout=60,
diffusion_batch_size=2,
async_chunk=False,
single_stage_id_filter=None,
omni_master_address="127.0.0.1",
@@ -928,7 +925,6 @@ class TestSingleStageReplicaInitialization:
model="fake-model",
config_path="/fake/stages.yaml",
stage_init_timeout=60,
diffusion_batch_size=2,
async_chunk=False,
single_stage_id_filter=None,
omni_master_address="127.0.0.1",
@@ -950,7 +946,6 @@ class TestSingleStageReplicaInitialization:
model="fake-model",
config_path="/fake/stages.yaml",
stage_init_timeout=60,
diffusion_batch_size=2,
async_chunk=False,
single_stage_id_filter=None,
omni_master_address="127.0.0.1",
@@ -1002,7 +997,6 @@ class TestSingleStageReplicaInitialization:
model="fake-model",
config_path="/fake/stages.yaml",
stage_init_timeout=60,
diffusion_batch_size=2,
async_chunk=False,
single_stage_id_filter=None,
omni_master_address="127.0.0.1",
@@ -1068,7 +1062,6 @@ class TestSingleStageReplicaInitialization:
model="fake-model",
config_path="/fake/stages.yaml",
stage_init_timeout=60,
diffusion_batch_size=4,
async_chunk=False,
single_stage_id_filter=None,
omni_master_address="127.0.0.1",
@@ -1120,7 +1113,6 @@ class TestSingleStageReplicaInitialization:
model="fake-model",
config_path="/fake/stages.yaml",
stage_init_timeout=60,
diffusion_batch_size=4,
async_chunk=False,
single_stage_id_filter=None,
omni_master_address="127.0.0.1",
@@ -1141,7 +1133,7 @@ class TestSingleStageReplicaInitialization:
runtime._init_visible_devices_baseline = "0"
mocker.patch.object(runtime_mod, "inject_kv_stage_info")
od_config = SimpleNamespace(max_num_seqs=None, parallel_config=SimpleNamespace(world_size=1))
od_config = SimpleNamespace(max_num_seqs=4, parallel_config=SimpleNamespace(world_size=1))
mocker.patch("vllm_omni.engine.stage_engine_startup.build_diffusion_config", return_value=od_config)
mock_register = mocker.patch(
"vllm_omni.engine.stage_engine_startup.register_stage_with_omni_master",
@@ -1178,7 +1170,7 @@ class TestSingleStageReplicaInitialization:
os.environ[device_env_var] = prev_device_env
assert result is sentinel_client
assert od_config.max_num_seqs is None
assert od_config.max_num_seqs == 4
mock_register.assert_called_once_with(
omni_master_address="127.0.0.1",
omni_master_port=25000,
@@ -1203,7 +1195,6 @@ class TestSingleStageReplicaInitialization:
request_address="tcp://127.0.0.1:26002",
response_address="tcp://127.0.0.1:26003",
proc_manager=mocker.ANY,
batch_size=4,
)
assert mock_from_addresses.call_args.kwargs["proc_manager"].proc is proc
@@ -1216,7 +1207,6 @@ class TestSingleStageReplicaInitialization:
model="fake-model",
config_path="/fake/stages.yaml",
stage_init_timeout=60,
diffusion_batch_size=4,
async_chunk=False,
single_stage_id_filter=None,
omni_master_address="127.0.0.1",
@@ -1258,7 +1248,6 @@ class TestSingleStageReplicaInitialization:
model="fake-model",
config_path="/fake/stages.yaml",
stage_init_timeout=60,
diffusion_batch_size=4,
async_chunk=False,
single_stage_id_filter=None,
omni_master_address="127.0.0.1",