[Refactor] Omni Stage Runtime and Distributed Replica Control Plane (#3855)
Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com> Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com> Co-authored-by: Chenguang ZHENG <645327136@qq.com>
This commit is contained in:
28 files changed
+3490
-3094
No files matched your search
@@ -1,21 +1,20 @@
|
||||
import concurrent.futures
|
||||
import contextlib
|
||||
import importlib
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
import types
|
||||
from contextlib import contextmanager
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec
|
||||
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine, StageRuntimeInfo
|
||||
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
|
||||
from vllm_omni.engine.stage_init_utils import (
|
||||
LogicalStageInitPlan,
|
||||
ReplicaInitPlan,
|
||||
build_stage0_input_processor,
|
||||
compute_replica_layout,
|
||||
)
|
||||
from vllm_omni.engine.stage_runtime import StageRuntime
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
@@ -63,7 +62,7 @@ def _make_diffusion_metadata(stage_id: int, *, replica_id: int = 0, final_output
|
||||
def _make_llm_plan(
|
||||
stage_idx: int,
|
||||
*,
|
||||
configured_stage_id: int,
|
||||
stage_id: int,
|
||||
vllm_config: object,
|
||||
num_replicas: int = 1,
|
||||
final_output: bool = False,
|
||||
@@ -73,7 +72,7 @@ def _make_llm_plan(
|
||||
replicas: list[ReplicaInitPlan] = []
|
||||
for replica_id in range(num_replicas):
|
||||
stage_cfg = types.SimpleNamespace(
|
||||
stage_id=configured_stage_id,
|
||||
stage_id=stage_id,
|
||||
stage_type="llm",
|
||||
runtime=types.SimpleNamespace(devices=str(replica_id)),
|
||||
engine_args={},
|
||||
@@ -85,7 +84,7 @@ def _make_llm_plan(
|
||||
launch_mode="local",
|
||||
stage_cfg=stage_cfg,
|
||||
metadata=_make_llm_metadata(
|
||||
configured_stage_id,
|
||||
stage_id,
|
||||
replica_id=replica_id,
|
||||
final_output=final_output,
|
||||
final_output_type=final_output_type,
|
||||
@@ -99,7 +98,7 @@ def _make_llm_plan(
|
||||
)
|
||||
return LogicalStageInitPlan(
|
||||
stage_idx=stage_idx,
|
||||
configured_stage_id=configured_stage_id,
|
||||
stage_id=stage_id,
|
||||
replicas=replicas,
|
||||
)
|
||||
|
||||
@@ -107,13 +106,13 @@ def _make_llm_plan(
|
||||
def _make_diffusion_plan(
|
||||
stage_idx: int,
|
||||
*,
|
||||
configured_stage_id: int,
|
||||
stage_id: int,
|
||||
num_replicas: int = 1,
|
||||
):
|
||||
replicas: list[ReplicaInitPlan] = []
|
||||
for replica_id in range(num_replicas):
|
||||
stage_cfg = types.SimpleNamespace(
|
||||
stage_id=configured_stage_id,
|
||||
stage_id=stage_id,
|
||||
stage_type="diffusion",
|
||||
runtime=types.SimpleNamespace(devices=str(replica_id)),
|
||||
engine_args={},
|
||||
@@ -124,14 +123,14 @@ def _make_diffusion_plan(
|
||||
num_replicas=num_replicas,
|
||||
launch_mode="local",
|
||||
stage_cfg=stage_cfg,
|
||||
metadata=_make_diffusion_metadata(configured_stage_id, replica_id=replica_id),
|
||||
metadata=_make_diffusion_metadata(stage_id, replica_id=replica_id),
|
||||
stage_connector_spec={},
|
||||
omni_kv_connector=(None, None, None),
|
||||
)
|
||||
)
|
||||
return LogicalStageInitPlan(
|
||||
stage_idx=stage_idx,
|
||||
configured_stage_id=configured_stage_id,
|
||||
stage_id=stage_id,
|
||||
replicas=replicas,
|
||||
)
|
||||
|
||||
@@ -165,7 +164,7 @@ def test_collect_initialized_clients_for_cleanup_deduplicates_clients():
|
||||
shared = types.SimpleNamespace(name="shared")
|
||||
extra = types.SimpleNamespace(name="extra")
|
||||
|
||||
cleanup_clients = AsyncOmniEngine._collect_initialized_clients_for_cleanup(
|
||||
cleanup_clients = StageRuntime._collect_initialized_clients_for_cleanup(
|
||||
stage_pools=[types.SimpleNamespace(clients=[shared, None])],
|
||||
initialized_clients_by_stage={0: [shared], 1: [extra]},
|
||||
)
|
||||
@@ -173,35 +172,36 @@ def test_collect_initialized_clients_for_cleanup_deduplicates_clients():
|
||||
assert cleanup_clients == [shared, extra]
|
||||
|
||||
|
||||
def test_initialize_diffusion_replica_restores_device_visibility_after_local_init(monkeypatch):
|
||||
import vllm_omni.engine.async_omni_engine as engine_mod
|
||||
def test_initialize_local_diffusion_replica_restores_device_visibility_after_local_init(monkeypatch):
|
||||
import vllm_omni.engine.stage_runtime as runtime_mod
|
||||
from vllm_omni.engine.stage_engine_startup import StageReplicaResources
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
engine.model = "dummy-model"
|
||||
engine.num_stages = 1
|
||||
engine.diffusion_batch_size = 1
|
||||
engine.single_stage_mode = False
|
||||
engine._omni_master_server = None
|
||||
engine.stage_configs = []
|
||||
runtime = StageRuntime(
|
||||
stage_configs=[],
|
||||
model="dummy-model",
|
||||
config_path="dummy-config",
|
||||
stage_init_timeout=1,
|
||||
diffusion_batch_size=1,
|
||||
async_chunk=False,
|
||||
)
|
||||
|
||||
plan = _make_diffusion_plan(0, configured_stage_id=0).replicas[0]
|
||||
plan = _make_diffusion_plan(0, stage_id=0).replicas[0]
|
||||
|
||||
env_var = current_omni_platform.device_control_env_var
|
||||
old_env = os.environ.get(env_var)
|
||||
os.environ[env_var] = "0,1"
|
||||
runtime._init_visible_devices_baseline = "0,1"
|
||||
|
||||
@contextmanager
|
||||
def _fake_stage_runtime_setup(_stage_id, _runtime_cfg):
|
||||
current_omni_platform.set_device_control_env_var("1")
|
||||
yield
|
||||
|
||||
monkeypatch.setattr(engine_mod, "stage_runtime_setup", _fake_stage_runtime_setup)
|
||||
monkeypatch.setattr(engine_mod, "inject_kv_stage_info", lambda *_: None)
|
||||
monkeypatch.setattr(engine_mod, "initialize_diffusion_stage", lambda *_, **__: types.SimpleNamespace())
|
||||
monkeypatch.setattr(runtime_mod, "inject_kv_stage_info", lambda *_: None)
|
||||
monkeypatch.setattr(
|
||||
runtime_mod,
|
||||
"launch_diffusion_stage_replica",
|
||||
lambda **_: (types.SimpleNamespace(), StageReplicaResources()),
|
||||
)
|
||||
|
||||
try:
|
||||
engine._initialize_diffusion_replica(plan, stage_init_timeout=1, stage_launch_lock=threading.Lock())
|
||||
runtime._initialize_local_diffusion_replica(plan, stage_init_timeout=1)
|
||||
assert os.environ.get(env_var) == "0,1"
|
||||
finally:
|
||||
if old_env is None:
|
||||
@@ -210,68 +210,63 @@ def test_initialize_diffusion_replica_restores_device_visibility_after_local_ini
|
||||
os.environ[env_var] = old_env
|
||||
|
||||
|
||||
def test_initialize_diffusion_replica_passes_stage_init_timeout_and_inline_flag(monkeypatch):
|
||||
import vllm_omni.engine.async_omni_engine as engine_mod
|
||||
def test_initialize_local_diffusion_replica_passes_stage_init_timeout_and_inline_flag(monkeypatch):
|
||||
import vllm_omni.engine.stage_runtime as runtime_mod
|
||||
from vllm_omni.engine.stage_engine_startup import StageReplicaResources
|
||||
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
engine.model = "dummy-model"
|
||||
engine.num_stages = 1
|
||||
engine.diffusion_batch_size = 4
|
||||
engine.single_stage_mode = False
|
||||
engine._omni_master_server = None
|
||||
engine.stage_configs = []
|
||||
runtime = StageRuntime(
|
||||
stage_configs=[types.SimpleNamespace()],
|
||||
model="dummy-model",
|
||||
config_path="dummy-config",
|
||||
stage_init_timeout=1,
|
||||
diffusion_batch_size=4,
|
||||
async_chunk=False,
|
||||
)
|
||||
|
||||
plan = _make_diffusion_plan(0, configured_stage_id=0).replicas[0]
|
||||
plan = _make_diffusion_plan(0, stage_id=0).replicas[0]
|
||||
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
@contextmanager
|
||||
def _noop_stage_runtime_setup(*_):
|
||||
yield
|
||||
monkeypatch.setattr(runtime_mod, "inject_kv_stage_info", lambda *_: None)
|
||||
|
||||
monkeypatch.setattr(engine_mod, "stage_runtime_setup", _noop_stage_runtime_setup)
|
||||
monkeypatch.setattr(engine_mod, "inject_kv_stage_info", lambda *_: None)
|
||||
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()
|
||||
|
||||
def _capture_initialize_diffusion_stage(
|
||||
stage_id, _model, _stage_cfg, _metadata, *, stage_init_timeout, batch_size, use_inline
|
||||
):
|
||||
captured["stage_id"] = stage_id
|
||||
captured["stage_init_timeout"] = stage_init_timeout
|
||||
captured["batch_size"] = batch_size
|
||||
captured["use_inline"] = use_inline
|
||||
return types.SimpleNamespace()
|
||||
monkeypatch.setattr(runtime_mod, "launch_diffusion_stage_replica", _capture_launch_diffusion_stage_replica)
|
||||
|
||||
monkeypatch.setattr(engine_mod, "initialize_diffusion_stage", _capture_initialize_diffusion_stage)
|
||||
|
||||
engine._initialize_diffusion_replica(plan, stage_init_timeout=302, stage_launch_lock=threading.Lock())
|
||||
runtime._initialize_local_diffusion_replica(plan, stage_init_timeout=302)
|
||||
|
||||
assert captured == {
|
||||
"stage_id": 0,
|
||||
"stage_init_timeout": 302,
|
||||
"batch_size": 4,
|
||||
"use_inline": True,
|
||||
"omni_master_server": None,
|
||||
}
|
||||
|
||||
|
||||
def test_initialize_stages_exposes_logical_stage_views_and_builds_top_level_input_processor(monkeypatch):
|
||||
import vllm_omni.engine.async_omni_engine as engine_mod
|
||||
def test_stage_runtime_initializes_stage_pools(monkeypatch):
|
||||
import vllm_omni.engine.stage_runtime as runtime_mod
|
||||
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
engine.model = "dummy-model"
|
||||
engine.config_path = "dummy-config"
|
||||
engine.num_stages = 2
|
||||
engine.async_chunk = False
|
||||
engine.diffusion_batch_size = 1
|
||||
engine.single_stage_mode = False
|
||||
engine._single_stage_id_filter = None
|
||||
engine._omni_master_server = None
|
||||
engine.stage_configs = [types.SimpleNamespace(), types.SimpleNamespace()]
|
||||
runtime = StageRuntime(
|
||||
stage_configs=[types.SimpleNamespace(), types.SimpleNamespace()],
|
||||
model="dummy-model",
|
||||
config_path="dummy-config",
|
||||
stage_init_timeout=1,
|
||||
diffusion_batch_size=1,
|
||||
async_chunk=False,
|
||||
)
|
||||
|
||||
cfg0 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64))
|
||||
cfg1 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64))
|
||||
stage_plans = [
|
||||
_make_llm_plan(0, configured_stage_id=0, vllm_config=cfg0, num_replicas=2, is_comprehension=True),
|
||||
_make_llm_plan(1, configured_stage_id=1, vllm_config=cfg1, final_output=True),
|
||||
_make_llm_plan(0, stage_id=0, vllm_config=cfg0, num_replicas=2, is_comprehension=True),
|
||||
_make_llm_plan(1, stage_id=1, vllm_config=cfg1, final_output=True),
|
||||
]
|
||||
|
||||
stage0_client_r0 = types.SimpleNamespace(
|
||||
@@ -302,51 +297,43 @@ def test_initialize_stages_exposes_logical_stage_views_and_builds_top_level_inpu
|
||||
|
||||
stage0_output_processor = object()
|
||||
stage1_output_processor = object()
|
||||
top_level_input_processor = object()
|
||||
|
||||
monkeypatch.setattr(engine_mod, "prepare_engine_environment", lambda: None)
|
||||
monkeypatch.setattr(engine_mod, "load_omni_transfer_config_for_model", lambda *_: None)
|
||||
monkeypatch.setattr(engine_mod, "compute_replica_layout", lambda _cfgs: ([2, 1], {}))
|
||||
monkeypatch.setattr(engine, "_build_logical_stage_init_plans", lambda *_: (stage_plans, None))
|
||||
monkeypatch.setattr(engine, "_initialize_stage_replicas", lambda *_: initialized_clients)
|
||||
monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: stage_plans)
|
||||
monkeypatch.setattr(runtime, "_initialize_stage_replicas", lambda *_: initialized_clients)
|
||||
monkeypatch.setattr(
|
||||
engine_mod,
|
||||
runtime_mod,
|
||||
"build_llm_stage_output_processor",
|
||||
lambda plan, _cfg, **_kw: stage0_output_processor if plan.stage_idx == 0 else stage1_output_processor,
|
||||
)
|
||||
monkeypatch.setattr(engine_mod, "build_stage0_input_processor", lambda _cfg: top_level_input_processor)
|
||||
|
||||
engine._initialize_stages(stage_init_timeout=1)
|
||||
runtime.initialize()
|
||||
|
||||
assert len(engine.stage_pools) == 2
|
||||
assert engine.input_processor is top_level_input_processor
|
||||
assert engine.stage_clients == [stage0_client_r0, stage1_client_r0]
|
||||
assert engine.stage_vllm_configs == [cfg0, cfg1]
|
||||
assert engine.output_processors == [stage0_output_processor, stage1_output_processor]
|
||||
assert engine.default_sampling_params_list == [
|
||||
stage0_client_r0.default_sampling_params,
|
||||
stage1_client_r0.default_sampling_params,
|
||||
]
|
||||
assert engine.stage_metadata == [
|
||||
StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm"),
|
||||
StageRuntimeInfo(final_output=True, final_output_type=None, stage_type="llm"),
|
||||
]
|
||||
assert len(runtime.stage_pools) == 2
|
||||
assert runtime.stage_pools[0].stage_client is stage0_client_r0
|
||||
assert runtime.stage_pools[1].stage_client is stage1_client_r0
|
||||
assert runtime.stage_pools[0].stage_vllm_config is cfg0
|
||||
assert runtime.stage_pools[1].stage_vllm_config is cfg1
|
||||
assert runtime.stage_pools[0].output_processor is stage0_output_processor
|
||||
assert runtime.stage_pools[1].output_processor is stage1_output_processor
|
||||
|
||||
|
||||
def test_build_logical_stage_init_plans_applies_replica_device_splits(monkeypatch):
|
||||
import vllm_omni.engine.async_omni_engine as engine_mod
|
||||
import vllm_omni.engine.stage_runtime as runtime_mod
|
||||
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
engine.model = "dummy-model"
|
||||
engine.async_chunk = False
|
||||
engine.single_stage_mode = False
|
||||
engine._single_stage_id_filter = None
|
||||
engine.stage_configs = [
|
||||
types.SimpleNamespace(stage_id=0, stage_type="llm", engine_args={}, runtime=types.SimpleNamespace(devices="0")),
|
||||
types.SimpleNamespace(
|
||||
stage_id=1, stage_type="llm", engine_args={}, runtime=types.SimpleNamespace(devices="1,2,3")
|
||||
),
|
||||
]
|
||||
runtime = StageRuntime(
|
||||
stage_configs=[
|
||||
types.SimpleNamespace(
|
||||
stage_id=0, stage_type="llm", engine_args={}, runtime=types.SimpleNamespace(devices="0")
|
||||
),
|
||||
types.SimpleNamespace(
|
||||
stage_id=1, stage_type="llm", engine_args={}, runtime=types.SimpleNamespace(devices="1,2,3")
|
||||
),
|
||||
],
|
||||
model="dummy-model",
|
||||
config_path="dummy-config",
|
||||
stage_init_timeout=1,
|
||||
diffusion_batch_size=1,
|
||||
async_chunk=False,
|
||||
)
|
||||
|
||||
metadata_by_stage = {
|
||||
0: _make_llm_metadata(0),
|
||||
@@ -354,78 +341,46 @@ def test_build_logical_stage_init_plans_applies_replica_device_splits(monkeypatc
|
||||
}
|
||||
|
||||
monkeypatch.setattr(
|
||||
engine_mod,
|
||||
runtime_mod,
|
||||
"extract_stage_metadata",
|
||||
lambda cfg: types.SimpleNamespace(**metadata_by_stage[cfg.stage_id].__dict__),
|
||||
)
|
||||
monkeypatch.setattr(engine_mod, "get_stage_connector_spec", lambda **_: {})
|
||||
monkeypatch.setattr(engine_mod, "resolve_omni_kv_config_for_stage", lambda *_: (None, None, None))
|
||||
monkeypatch.setattr(engine_mod, "build_engine_args_dict", lambda *_, **__: {})
|
||||
monkeypatch.setattr(runtime_mod, "get_stage_connector_spec", lambda **_: {})
|
||||
monkeypatch.setattr(runtime_mod, "resolve_omni_kv_config_for_stage", lambda *_: (None, None, None))
|
||||
monkeypatch.setattr(runtime_mod, "build_engine_args_dict", lambda *_, **__: {})
|
||||
monkeypatch.setattr(
|
||||
engine_mod,
|
||||
runtime_mod,
|
||||
"build_vllm_config",
|
||||
lambda stage_cfg, *_args, **_kwargs: (types.SimpleNamespace(tag=f"cfg-{stage_cfg.stage_id}"), object),
|
||||
)
|
||||
|
||||
stage_plans, prompt_expand_func = engine._build_logical_stage_init_plans(
|
||||
stage_plans = runtime._build_logical_stage_init_plans(
|
||||
omni_transfer_config=None,
|
||||
replicas_per_stage=[1, 3],
|
||||
replica_devices_map={1: ["1", "2", "3"]},
|
||||
)
|
||||
|
||||
assert prompt_expand_func is None
|
||||
assert [plan.configured_stage_id for plan in stage_plans] == [0, 1]
|
||||
assert [plan.stage_id for plan in stage_plans] == [0, 1]
|
||||
assert [replica.stage_cfg.runtime.devices for replica in stage_plans[1].replicas] == ["1", "2", "3"]
|
||||
assert [replica.replica_id for replica in stage_plans[1].replicas] == [0, 1, 2]
|
||||
assert all(replica.num_replicas == 3 for replica in stage_plans[1].replicas)
|
||||
|
||||
|
||||
def test_initialize_stage_replicas_reuses_long_lived_executor(monkeypatch):
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
engine._stage_init_executor = None
|
||||
|
||||
cfg = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64))
|
||||
stage_plans = [_make_llm_plan(0, configured_stage_id=0, vllm_config=cfg, num_replicas=1)]
|
||||
|
||||
created_executors = []
|
||||
|
||||
class FakeExecutor:
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.shutdown_calls = []
|
||||
created_executors.append(self)
|
||||
|
||||
def submit(self, fn, *args, **kwargs):
|
||||
future = concurrent.futures.Future()
|
||||
future.set_result(fn(*args, **kwargs))
|
||||
return future
|
||||
|
||||
def shutdown(self, wait=False, cancel_futures=False):
|
||||
self.shutdown_calls.append((wait, cancel_futures))
|
||||
|
||||
monkeypatch.setattr(concurrent.futures, "ThreadPoolExecutor", FakeExecutor)
|
||||
monkeypatch.setattr(
|
||||
engine,
|
||||
"_initialize_replica",
|
||||
lambda plan, _stage_init_timeout, _stage_launch_lock: f"client-{plan.replica_id}",
|
||||
)
|
||||
|
||||
first = engine._initialize_stage_replicas(stage_plans, stage_init_timeout=1)
|
||||
second = engine._initialize_stage_replicas(stage_plans, stage_init_timeout=1)
|
||||
|
||||
assert first == {0: ["client-0"]}
|
||||
assert second == {0: ["client-0"]}
|
||||
assert len(created_executors) == 1
|
||||
assert engine._stage_init_executor is created_executors[0]
|
||||
|
||||
|
||||
def test_initialize_stage_replicas_collects_results_by_stage_and_replica_id(monkeypatch):
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
runtime = StageRuntime(
|
||||
stage_configs=[],
|
||||
model="dummy-model",
|
||||
config_path="dummy-config",
|
||||
stage_init_timeout=123,
|
||||
diffusion_batch_size=1,
|
||||
async_chunk=False,
|
||||
)
|
||||
|
||||
cfg0 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64))
|
||||
cfg1 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64))
|
||||
stage_plans = [
|
||||
_make_llm_plan(0, configured_stage_id=0, vllm_config=cfg0, num_replicas=2),
|
||||
_make_llm_plan(1, configured_stage_id=1, vllm_config=cfg1, num_replicas=2),
|
||||
_make_llm_plan(0, stage_id=0, vllm_config=cfg0, num_replicas=2),
|
||||
_make_llm_plan(1, stage_id=1, vllm_config=cfg1, num_replicas=2),
|
||||
]
|
||||
|
||||
clients = {
|
||||
@@ -435,13 +390,13 @@ def test_initialize_stage_replicas_collects_results_by_stage_and_replica_id(monk
|
||||
(1, 1): types.SimpleNamespace(name="stage1-replica1"),
|
||||
}
|
||||
|
||||
def _initialize_replica(plan, _stage_init_timeout, _stage_launch_lock):
|
||||
def _initialize_replica(plan, _stage_init_timeout):
|
||||
time.sleep(0.02 * (3 - plan.metadata.stage_id - plan.replica_id))
|
||||
return clients[(plan.metadata.stage_id, plan.replica_id)]
|
||||
|
||||
monkeypatch.setattr(engine, "_initialize_replica", _initialize_replica)
|
||||
monkeypatch.setattr(runtime, "_initialize_replica", _initialize_replica)
|
||||
|
||||
initialized_clients = engine._initialize_stage_replicas(stage_plans, stage_init_timeout=123)
|
||||
initialized_clients = runtime._initialize_stage_replicas(stage_plans, stage_init_timeout=123)
|
||||
|
||||
assert initialized_clients == {
|
||||
0: [clients[(0, 0)], clients[(0, 1)]],
|
||||
@@ -449,148 +404,120 @@ def test_initialize_stage_replicas_collects_results_by_stage_and_replica_id(monk
|
||||
}
|
||||
|
||||
|
||||
def test_shutdown_closes_long_lived_stage_init_executor():
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
engine._shutdown_called = False
|
||||
engine._weak_finalizer = None
|
||||
engine._omni_master_server = None
|
||||
engine._coordinator_runtime = None
|
||||
engine.orchestrator_thread = types.SimpleNamespace(join=lambda timeout: None, is_alive=lambda: False)
|
||||
def test_remote_replicas_use_distinct_init_group_keys():
|
||||
runtime = StageRuntime(
|
||||
stage_configs=[],
|
||||
model="dummy-model",
|
||||
config_path="dummy-config",
|
||||
stage_init_timeout=123,
|
||||
diffusion_batch_size=1,
|
||||
async_chunk=False,
|
||||
)
|
||||
plan = _make_llm_plan(
|
||||
1,
|
||||
stage_id=1,
|
||||
vllm_config=types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64)),
|
||||
num_replicas=2,
|
||||
)
|
||||
|
||||
class FakeQueue:
|
||||
def __init__(self):
|
||||
self.closed = False
|
||||
self.sync_q = types.SimpleNamespace(put_nowait=lambda _msg: None)
|
||||
for replica in plan.replicas:
|
||||
replica.launch_mode = "remote"
|
||||
replica.metadata.runtime_cfg = None
|
||||
|
||||
def close(self):
|
||||
self.closed = True
|
||||
|
||||
class FakeExecutor:
|
||||
def __init__(self):
|
||||
self.shutdown_calls = []
|
||||
|
||||
def shutdown(self, wait=False, cancel_futures=False):
|
||||
self.shutdown_calls.append((wait, cancel_futures))
|
||||
|
||||
engine.request_queue = FakeQueue()
|
||||
engine.output_queue = FakeQueue()
|
||||
engine.rpc_output_queue = FakeQueue()
|
||||
engine._stage_init_executor = FakeExecutor()
|
||||
|
||||
engine.shutdown()
|
||||
|
||||
assert engine._stage_init_executor is None
|
||||
assert engine.request_queue.closed is True
|
||||
assert engine.output_queue.closed is True
|
||||
assert engine.rpc_output_queue.closed is True
|
||||
assert [runtime._replica_init_group_key(replica) for replica in plan.replicas] == [
|
||||
"remote:1:0",
|
||||
"remote:1:1",
|
||||
]
|
||||
|
||||
|
||||
def test_initialize_stages_cleans_up_successful_replicas_after_partial_multi_replica_failure(monkeypatch):
|
||||
import vllm_omni.engine.async_omni_engine as engine_mod
|
||||
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
engine.model = "dummy-model"
|
||||
engine.config_path = "dummy-config"
|
||||
engine.num_stages = 1
|
||||
engine.async_chunk = False
|
||||
engine.diffusion_batch_size = 1
|
||||
engine.single_stage_mode = False
|
||||
engine._single_stage_id_filter = None
|
||||
engine._omni_master_server = None
|
||||
engine._coordinator_runtime = None
|
||||
engine.stage_configs = [types.SimpleNamespace()]
|
||||
runtime = StageRuntime(
|
||||
stage_configs=[types.SimpleNamespace()],
|
||||
model="dummy-model",
|
||||
config_path="dummy-config",
|
||||
stage_init_timeout=1,
|
||||
diffusion_batch_size=1,
|
||||
async_chunk=False,
|
||||
)
|
||||
|
||||
cfg0 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64))
|
||||
stage_plans = [_make_llm_plan(0, configured_stage_id=0, vllm_config=cfg0, num_replicas=2)]
|
||||
stage_plans = [_make_llm_plan(0, stage_id=0, vllm_config=cfg0, num_replicas=2)]
|
||||
initialized_client = types.SimpleNamespace(shutdown=lambda: None)
|
||||
|
||||
monkeypatch.setattr(engine_mod, "prepare_engine_environment", lambda: None)
|
||||
monkeypatch.setattr(engine_mod, "load_omni_transfer_config_for_model", lambda *_: None)
|
||||
monkeypatch.setattr(engine_mod, "compute_replica_layout", lambda _cfgs: ([2], {}))
|
||||
monkeypatch.setattr(engine, "_build_logical_stage_init_plans", lambda *_: (stage_plans, None))
|
||||
monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: stage_plans)
|
||||
|
||||
def _initialize_replica(plan, _stage_init_timeout, _stage_launch_lock):
|
||||
def _initialize_replica(plan, _stage_init_timeout):
|
||||
if plan.replica_id == 0:
|
||||
return initialized_client
|
||||
time.sleep(0.05)
|
||||
raise RuntimeError("replica launch failed")
|
||||
|
||||
monkeypatch.setattr(engine, "_initialize_replica", _initialize_replica)
|
||||
monkeypatch.setattr(runtime, "_initialize_replica", _initialize_replica)
|
||||
|
||||
captured_cleanup: list[list[object]] = []
|
||||
|
||||
def _capture_shutdown(clients):
|
||||
captured_cleanup.append(list(clients))
|
||||
|
||||
monkeypatch.setattr(engine, "_shutdown_initialized_clients", _capture_shutdown)
|
||||
monkeypatch.setattr(runtime, "_shutdown_initialized_clients", _capture_shutdown)
|
||||
|
||||
with pytest.raises(RuntimeError, match="replica launch failed"):
|
||||
engine._initialize_stages(stage_init_timeout=1)
|
||||
runtime.initialize()
|
||||
|
||||
assert captured_cleanup == [[initialized_client]]
|
||||
|
||||
|
||||
def test_initialize_stages_cleans_up_late_successful_replicas_after_early_multi_replica_failure(monkeypatch):
|
||||
import vllm_omni.engine.async_omni_engine as engine_mod
|
||||
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
engine.model = "dummy-model"
|
||||
engine.config_path = "dummy-config"
|
||||
engine.num_stages = 1
|
||||
engine.async_chunk = False
|
||||
engine.diffusion_batch_size = 1
|
||||
engine.single_stage_mode = False
|
||||
engine._single_stage_id_filter = None
|
||||
engine._omni_master_server = None
|
||||
engine._coordinator_runtime = None
|
||||
engine.stage_configs = [types.SimpleNamespace()]
|
||||
runtime = StageRuntime(
|
||||
stage_configs=[types.SimpleNamespace()],
|
||||
model="dummy-model",
|
||||
config_path="dummy-config",
|
||||
stage_init_timeout=1,
|
||||
diffusion_batch_size=1,
|
||||
async_chunk=False,
|
||||
)
|
||||
|
||||
cfg0 = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64))
|
||||
stage_plans = [_make_llm_plan(0, configured_stage_id=0, vllm_config=cfg0, num_replicas=2)]
|
||||
stage_plans = [_make_llm_plan(0, stage_id=0, vllm_config=cfg0, num_replicas=2)]
|
||||
initialized_client = types.SimpleNamespace(shutdown=lambda: None)
|
||||
replica1_started = threading.Event()
|
||||
|
||||
monkeypatch.setattr(engine_mod, "prepare_engine_environment", lambda: None)
|
||||
monkeypatch.setattr(engine_mod, "load_omni_transfer_config_for_model", lambda *_: None)
|
||||
monkeypatch.setattr(engine_mod, "compute_replica_layout", lambda _cfgs: ([2], {}))
|
||||
monkeypatch.setattr(engine, "_build_logical_stage_init_plans", lambda *_: (stage_plans, None))
|
||||
monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: stage_plans)
|
||||
|
||||
def _initialize_replica(plan, _stage_init_timeout, _stage_launch_lock):
|
||||
if plan.replica_id == 0:
|
||||
assert replica1_started.wait(timeout=1)
|
||||
raise RuntimeError("replica launch failed")
|
||||
replica1_started.set()
|
||||
time.sleep(0.05)
|
||||
return initialized_client
|
||||
def _initialize_stage_replicas(_stage_plans, _stage_init_timeout):
|
||||
exc = RuntimeError("replica launch failed")
|
||||
exc._initialized_clients_by_stage = {0: [None, initialized_client]}
|
||||
raise exc
|
||||
|
||||
monkeypatch.setattr(engine, "_initialize_replica", _initialize_replica)
|
||||
monkeypatch.setattr(runtime, "_initialize_stage_replicas", _initialize_stage_replicas)
|
||||
|
||||
captured_cleanup: list[list[object]] = []
|
||||
|
||||
def _capture_shutdown(clients):
|
||||
captured_cleanup.append(list(clients))
|
||||
|
||||
monkeypatch.setattr(engine, "_shutdown_initialized_clients", _capture_shutdown)
|
||||
monkeypatch.setattr(runtime, "_shutdown_initialized_clients", _capture_shutdown)
|
||||
|
||||
with pytest.raises(RuntimeError, match="replica launch failed"):
|
||||
engine._initialize_stages(stage_init_timeout=1)
|
||||
runtime.initialize()
|
||||
|
||||
assert captured_cleanup == [[initialized_client]]
|
||||
|
||||
|
||||
def test_initialize_llm_replica_passes_stage_init_timeout_to_complete_stage_handshake(monkeypatch):
|
||||
import vllm_omni.engine.async_omni_engine as engine_mod
|
||||
def test_initialize_local_llm_replica_passes_stage_init_timeout_to_complete_stage_handshake(monkeypatch):
|
||||
import vllm_omni.engine.stage_runtime as runtime_mod
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
engine.model = "dummy-model"
|
||||
engine.single_stage_mode = False
|
||||
engine._omni_master_server = None
|
||||
engine.stage_configs = []
|
||||
runtime = StageRuntime(
|
||||
stage_configs=[],
|
||||
model="dummy-model",
|
||||
config_path="dummy-config",
|
||||
stage_init_timeout=302,
|
||||
diffusion_batch_size=1,
|
||||
async_chunk=False,
|
||||
)
|
||||
|
||||
fake_vllm_config = types.SimpleNamespace()
|
||||
fake_addresses = types.SimpleNamespace(inputs=["in"], outputs=["out"], frontend_stats_publish_address=None)
|
||||
fake_proc = types.SimpleNamespace()
|
||||
captured_timeout: int | None = None
|
||||
|
||||
plan = ReplicaInitPlan(
|
||||
@@ -603,34 +530,38 @@ def test_initialize_llm_replica_passes_stage_init_timeout_to_complete_stage_hand
|
||||
omni_kv_connector=(None, None, None),
|
||||
stage_vllm_config=fake_vllm_config,
|
||||
executor_class=object,
|
||||
engine_args_dict={},
|
||||
)
|
||||
|
||||
device_env_var = current_omni_platform.device_control_env_var
|
||||
prev_device_env = os.environ.get(device_env_var)
|
||||
os.environ[device_env_var] = "0"
|
||||
|
||||
@contextmanager
|
||||
def _noop_stage_runtime_setup(*_):
|
||||
yield
|
||||
|
||||
monkeypatch.setattr(engine_mod, "stage_runtime_setup", _noop_stage_runtime_setup)
|
||||
monkeypatch.setattr(engine_mod, "build_engine_args_dict", lambda *_, **__: {})
|
||||
monkeypatch.setattr(engine_mod, "acquire_device_locks", lambda *_: [])
|
||||
monkeypatch.setattr(engine_mod, "spawn_stage_core", lambda **_: (fake_addresses, fake_proc, "ipc://handshake"))
|
||||
|
||||
def _capture_stage_timeout(_proc, _handshake_addr, _addresses, _vllm_cfg, handshake_timeout):
|
||||
def _capture_acquire_device_locks(*_args):
|
||||
nonlocal captured_timeout
|
||||
captured_timeout = handshake_timeout
|
||||
captured_timeout = _args[2]
|
||||
return []
|
||||
|
||||
monkeypatch.setattr(engine_mod, "complete_stage_handshake", _capture_stage_timeout)
|
||||
monkeypatch.setattr(runtime_mod, "acquire_device_locks", _capture_acquire_device_locks)
|
||||
|
||||
from vllm_omni.engine.stage_engine_startup import StageReplicaResources
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _fake_launch_stage_replica(**_kwargs):
|
||||
yield StageReplicaResources(
|
||||
manager=types.SimpleNamespace(shutdown=lambda: None),
|
||||
addresses=fake_addresses,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(runtime_mod, "launch_stage_replica", _fake_launch_stage_replica)
|
||||
monkeypatch.setattr(
|
||||
engine_mod.StageEngineCoreClientBase,
|
||||
runtime_mod.StageEngineCoreClientBase,
|
||||
"make_async_mp_client",
|
||||
staticmethod(lambda **_: types.SimpleNamespace(shutdown=lambda: None)),
|
||||
)
|
||||
|
||||
try:
|
||||
engine._initialize_llm_replica(plan, 302, threading.Lock())
|
||||
runtime._initialize_local_llm_replica(plan, 302)
|
||||
finally:
|
||||
if prev_device_env is None:
|
||||
os.environ.pop(device_env_var, None)
|
||||
|
||||
@@ -0,0 +1,158 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import threading
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.distributed.omni_coordinator.messages import ReplicaStatus
|
||||
from vllm_omni.engine.membership_controller import MembershipController
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class FakeHub:
|
||||
def __init__(self, snapshots=None):
|
||||
self.snapshots = list(snapshots or [])
|
||||
self.closed = False
|
||||
|
||||
def get_replica_list(self):
|
||||
if self.snapshots:
|
||||
return self.snapshots.pop(0)
|
||||
return SimpleNamespace(replicas=[])
|
||||
|
||||
def close(self):
|
||||
self.closed = True
|
||||
|
||||
|
||||
class FakePool:
|
||||
def __init__(self, stage_id: int):
|
||||
self.stage_id = stage_id
|
||||
self.clients = []
|
||||
self.added = []
|
||||
self.removed = []
|
||||
self.invalidated = []
|
||||
self.hub = None
|
||||
self.lb = None
|
||||
|
||||
def attach_hub(self, hub):
|
||||
self.hub = hub
|
||||
|
||||
def attach_load_balancer(self, lb):
|
||||
self.lb = lb
|
||||
|
||||
def add_client(self, input_addr, client):
|
||||
self.added.append((input_addr, client))
|
||||
self.clients.append(client)
|
||||
return len(self.clients) - 1
|
||||
|
||||
def invalidate_addr(self, input_addr):
|
||||
self.invalidated.append(input_addr)
|
||||
return ["req-1", "req-2"]
|
||||
|
||||
def remove_client(self, input_addr):
|
||||
self.removed.append(input_addr)
|
||||
return SimpleNamespace(shutdown=lambda: self.removed.append("shutdown"))
|
||||
|
||||
|
||||
def _snapshot(*replicas):
|
||||
return SimpleNamespace(replicas=list(replicas))
|
||||
|
||||
|
||||
def _replica(stage_id: int, input_addr: str, status=ReplicaStatus.UP):
|
||||
return SimpleNamespace(stage_id=stage_id, input_addr=input_addr, status=status)
|
||||
|
||||
|
||||
def _controller(monkeypatch, pool, hub, remote_replica_factory=None):
|
||||
import vllm_omni.engine.membership_controller as membership_mod
|
||||
|
||||
monkeypatch.setattr(membership_mod, "OmniCoordClientForHub", lambda _addr: hub)
|
||||
if remote_replica_factory is None:
|
||||
|
||||
def remote_replica_factory(stage_id, replica_id):
|
||||
return SimpleNamespace(client_addresses={"input_address": f"tcp://stage-{stage_id}-replica-{replica_id}"})
|
||||
|
||||
return MembershipController(
|
||||
stage_pools=[pool],
|
||||
coordinator_pub_address="tcp://127.0.0.1:12345",
|
||||
load_balancer_factory=lambda: object(),
|
||||
remote_replica_factory=remote_replica_factory,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_watch_replica_list_unregisters_disappeared_replicas(monkeypatch):
|
||||
pool = FakePool(stage_id=0)
|
||||
hub = FakeHub(
|
||||
snapshots=[
|
||||
_snapshot(_replica(0, "tcp://gone")),
|
||||
_snapshot(),
|
||||
]
|
||||
)
|
||||
controller = _controller(monkeypatch, pool, hub)
|
||||
controller.WATCH_INTERVAL_S = 0
|
||||
unregistered = []
|
||||
|
||||
async def _handle_unregister(stage_id, input_addr, *args, **kwargs):
|
||||
unregistered.append((stage_id, input_addr))
|
||||
controller._shutdown_event.set()
|
||||
|
||||
controller.handle_unregister = _handle_unregister # type: ignore[method-assign]
|
||||
|
||||
await asyncio.wait_for(controller._watch_replica_list(), timeout=1)
|
||||
await controller.drain_tasks(timeout=1)
|
||||
|
||||
assert unregistered == [(0, "tcp://gone")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_shutdown_closes_hub_then_cancels_watcher(monkeypatch):
|
||||
pool = FakePool(stage_id=0)
|
||||
hub = FakeHub()
|
||||
controller = _controller(monkeypatch, pool, hub)
|
||||
watcher = asyncio.create_task(asyncio.sleep(10))
|
||||
controller._watcher_task = watcher
|
||||
|
||||
controller.shutdown()
|
||||
await asyncio.sleep(0)
|
||||
|
||||
assert hub.closed is True
|
||||
assert watcher.cancelled()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_drain_tasks_waits_for_membership_tasks(monkeypatch):
|
||||
pool = FakePool(stage_id=0)
|
||||
controller = _controller(monkeypatch, pool, FakeHub())
|
||||
completed = []
|
||||
|
||||
async def _task():
|
||||
await asyncio.sleep(0)
|
||||
completed.append(True)
|
||||
|
||||
controller._spawn_task(_task(), label="unit")
|
||||
await controller.drain_tasks(timeout=1)
|
||||
|
||||
assert completed == [True]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_do_register_offloads_remote_factory(monkeypatch):
|
||||
pool = FakePool(stage_id=0)
|
||||
factory_thread_ids = []
|
||||
|
||||
def _factory(stage_id, replica_id):
|
||||
factory_thread_ids.append(threading.get_ident())
|
||||
return SimpleNamespace(client_addresses={"input_address": f"tcp://stage-{stage_id}-replica-{replica_id}"})
|
||||
|
||||
controller = _controller(monkeypatch, pool, FakeHub(), remote_replica_factory=_factory)
|
||||
|
||||
await controller._do_register(0, 1)
|
||||
|
||||
assert factory_thread_ids
|
||||
assert factory_thread_ids[0] != threading.get_ident()
|
||||
assert pool.added[0][0] == "tcp://stage-0-replica-1"
|
||||
@@ -11,11 +11,9 @@ from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import janus
|
||||
import psutil
|
||||
import pytest
|
||||
from vllm.outputs import CompletionOutput, RequestOutput
|
||||
from vllm.sampling_params import SamplingParams
|
||||
from vllm.v1.engine.core_client import AsyncMPClient
|
||||
|
||||
from vllm_omni.engine.messages import (
|
||||
AbortRequestMessage,
|
||||
@@ -27,7 +25,6 @@ from vllm_omni.engine.messages import (
|
||||
StageSubmissionMessage,
|
||||
)
|
||||
from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState
|
||||
from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient
|
||||
from vllm_omni.engine.stage_pool import StagePool
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
@@ -180,42 +177,6 @@ class FakeOutputProcessor:
|
||||
return None
|
||||
|
||||
|
||||
class _FakeProc:
|
||||
pid = 1234
|
||||
|
||||
def __init__(self):
|
||||
self.terminated = False
|
||||
self.killed = False
|
||||
self.join_calls = []
|
||||
|
||||
def is_alive(self):
|
||||
return not self.terminated and not self.killed
|
||||
|
||||
def terminate(self):
|
||||
self.terminated = True
|
||||
|
||||
def kill(self):
|
||||
self.killed = True
|
||||
|
||||
def join(self, timeout=None):
|
||||
self.join_calls.append(timeout)
|
||||
|
||||
|
||||
class _FakeChildProc:
|
||||
def __init__(self):
|
||||
self.terminated = False
|
||||
self.killed = False
|
||||
|
||||
def is_running(self):
|
||||
return not self.terminated and not self.killed
|
||||
|
||||
def terminate(self):
|
||||
self.terminated = True
|
||||
|
||||
def kill(self):
|
||||
self.killed = True
|
||||
|
||||
|
||||
def _sampling_params(max_tokens: int = 4) -> SamplingParams:
|
||||
return SamplingParams(max_tokens=max_tokens)
|
||||
|
||||
@@ -437,61 +398,6 @@ async def _enqueue_abort_request(orchestrator_fixture: OrchestratorFixture, requ
|
||||
orchestrator_fixture.request_sync_q.put_nowait(AbortRequestMessage(request_ids=request_ids))
|
||||
|
||||
|
||||
def test_stage_engine_core_client_shutdown_cleans_children_if_base_shutdown_fails(monkeypatch):
|
||||
fake_proc = _FakeProc()
|
||||
fake_child = _FakeChildProc()
|
||||
|
||||
class FakePsutilProcess:
|
||||
def __init__(self, pid):
|
||||
assert pid == fake_proc.pid
|
||||
|
||||
def children(self, recursive=True):
|
||||
assert recursive
|
||||
return [fake_child]
|
||||
|
||||
def fail_base_shutdown(self, **kwargs):
|
||||
raise RuntimeError("base shutdown failed")
|
||||
|
||||
monkeypatch.setattr(psutil, "Process", FakePsutilProcess)
|
||||
monkeypatch.setattr(psutil, "wait_procs", lambda procs, timeout: (list(procs), []))
|
||||
monkeypatch.setattr(AsyncMPClient, "shutdown", fail_base_shutdown)
|
||||
|
||||
client = object.__new__(StageEngineCoreClient)
|
||||
client._proc = fake_proc
|
||||
|
||||
with pytest.raises(RuntimeError, match="base shutdown failed"):
|
||||
client.shutdown()
|
||||
|
||||
assert fake_proc.terminated
|
||||
assert fake_proc.join_calls == [5]
|
||||
assert fake_child.terminated
|
||||
|
||||
|
||||
def test_stage_engine_core_client_shutdown_kills_stubborn_children(monkeypatch):
|
||||
fake_proc = _FakeProc()
|
||||
fake_child = _FakeChildProc()
|
||||
|
||||
class FakePsutilProcess:
|
||||
def __init__(self, pid):
|
||||
assert pid == fake_proc.pid
|
||||
|
||||
def children(self, recursive=True):
|
||||
assert recursive
|
||||
return [fake_child]
|
||||
|
||||
monkeypatch.setattr(psutil, "Process", FakePsutilProcess)
|
||||
monkeypatch.setattr(psutil, "wait_procs", lambda procs, timeout: ([], list(procs)))
|
||||
monkeypatch.setattr(AsyncMPClient, "shutdown", lambda self, **kwargs: None)
|
||||
|
||||
client = object.__new__(StageEngineCoreClient)
|
||||
client._proc = fake_proc
|
||||
|
||||
client.shutdown()
|
||||
|
||||
assert fake_child.terminated
|
||||
assert fake_child.killed
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def orchestrator_factory():
|
||||
fixtures: list[OrchestratorFixture] = []
|
||||
|
||||
File diff suppressed because it is too large.
Load diff
@@ -1,15 +1,10 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for StageEngineCoreClient.check_health().
|
||||
|
||||
Uses object.__new__ to construct a minimal client — check_health only
|
||||
touches self.resources, self.stage_id, and self._proc.
|
||||
"""
|
||||
"""Tests for StageEngineCoreClient.check_health()."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from vllm.v1.engine.exceptions import EngineDeadError
|
||||
@@ -19,28 +14,19 @@ from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _make_client(*, engine_dead=False, proc_alive=True):
|
||||
def _make_client(*, engine_dead=False):
|
||||
client = object.__new__(StageEngineCoreClient)
|
||||
client.stage_id = 0
|
||||
client.resources = SimpleNamespace(engine_dead=engine_dead)
|
||||
client._proc = MagicMock(is_alive=MagicMock(return_value=proc_alive), exitcode=1)
|
||||
return client
|
||||
|
||||
|
||||
def test_check_health_passes_when_alive():
|
||||
client = _make_client(engine_dead=False, proc_alive=True)
|
||||
client = _make_client(engine_dead=False)
|
||||
client.check_health() # no exception
|
||||
|
||||
|
||||
def test_check_health_raises_when_resources_engine_dead():
|
||||
client = _make_client(engine_dead=True, proc_alive=True)
|
||||
client = _make_client(engine_dead=True)
|
||||
with pytest.raises(EngineDeadError, match="engine core is dead"):
|
||||
client.check_health()
|
||||
|
||||
|
||||
def test_check_health_raises_when_proc_not_alive():
|
||||
client = _make_client(engine_dead=False, proc_alive=False)
|
||||
with pytest.raises(EngineDeadError, match="not alive"):
|
||||
client.check_health()
|
||||
# Verify it set resources.engine_dead as a side effect
|
||||
assert client.resources.engine_dead is True
|
||||
Reference in new issue
Block a user