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

[Bugfix][Diffusion] Clean up inline worker lifecycle (#7986)

Browse Source 
Signed-off-by: andyluo7 <andy.luo@amd.com>
This commit is contained in:
andyluo7 authored and GitHub committed 2026-09-24 21:43:25 +08:00
1 parent 0a0e1884f7
commit 9b08ac8203
2 files changed
+278 -97
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/diffusion/test_worker_wrapper_base.py
+138 -17
View File
@@ -79,6 +79,33 @@ class MockCustomPipeline:
class TestWorkerWrapperBaseInitialization:
"""Test WorkerWrapperBase initialization behavior."""
def test_worker_init_failure_after_device_setup_shuts_down(self, mocker: MockerFixture):
"""A post-device constructor failure must unwind distributed state."""
class FakeRunner:
def __init__(self, **kwargs):
self.kwargs = kwargs
config = SimpleNamespace(
stage_id=0,
diffusion_model_runner_cls="tests.FakeRunner",
diffusion_load_format="default",
)
init_device = mocker.patch.object(DiffusionWorker, "init_device")
mocker.patch(
"vllm_omni.diffusion.worker.diffusion_worker.resolve_obj_by_qualname",
return_value=FakeRunner,
)
mocker.patch.object(DiffusionWorker, "_create_profiler", return_value=None)
mocker.patch.object(DiffusionWorker, "load_model", side_effect=RuntimeError("load failed"))
shutdown = mocker.patch.object(DiffusionWorker, "shutdown")
with pytest.raises(RuntimeError, match="load failed"):
DiffusionWorker(local_rank=0, rank=0, od_config=config)
init_device.assert_called_once_with()
shutdown.assert_called_once_with()
def test_basic_initialization(self, mocker: MockerFixture, mock_od_config):
"""Test basic initialization without extensions."""
mock_worker_init = mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
@@ -125,6 +152,35 @@ class TestWorkerWrapperBaseInitialization:
skip_load_model=False,
)
def test_custom_pipeline_init_failure_shuts_down_worker(self, mocker: MockerFixture, mock_od_config):
"""A failed custom reload must not abandon an initialized worker."""
instances = []
class FailingWorker:
def __init__(self, **kwargs):
self.kwargs = kwargs
self.shutdown_calls = 0
instances.append(self)
def re_init_pipeline(self, custom_pipeline_args):
raise RuntimeError(f"failed to load {custom_pipeline_args['pipeline_class']}")
def shutdown(self):
self.shutdown_calls += 1
mocker.patch.object(WorkerWrapperBase, "_prepare_worker_class", return_value=FailingWorker)
with pytest.raises(RuntimeError, match="failed to load test.Pipeline"):
WorkerWrapperBase(
gpu_id=0,
od_config=mock_od_config,
base_worker_class=DiffusionWorker,
custom_pipeline_args={"pipeline_class": "test.Pipeline"},
)
assert len(instances) == 1
assert instances[0].shutdown_calls == 1
# -------------------------------------------------------------------------
# Tests: Worker Extension Functionality
@@ -231,28 +287,101 @@ class TestWorkerWrapperBaseDelegation:
assert result is None
@pytest.mark.parametrize("native_kv", [False, True])
def test_worker_shutdown_disables_offloader_before_distributed_teardown(self, mocker: MockerFixture, native_kv):
def test_worker_shutdown_releases_model_before_cumem_pools(self, mocker: MockerFixture, native_kv):
events: list[str] = []
offload_backend = mocker.Mock()
offload_backend.disable.side_effect = lambda: events.append("offload")
kv_manager = mocker.Mock()
kv_manager.close.side_effect = lambda: events.append("kv")
destroy = mocker.patch(
"vllm_omni.diffusion.worker.diffusion_worker.destroy_distributed_env",
side_effect=lambda: events.append("distributed"),
)
worker = DiffusionWorker.__new__(DiffusionWorker)
worker.od_config = SimpleNamespace(enable_sleep_mode=True)
worker.model_runner = SimpleNamespace(
offload_backend=offload_backend,
kv_transfer_manager=None if native_kv else kv_manager,
_kv_transfer_manager=kv_manager,
)
worker.lora_manager = object()
worker.profiler = object()
worker._sleep_saved_buffers = {"buffer": object()}
worker._step_lora_state = {"request": (None, 1.0)}
worker._shutdown_complete = False
worker._owns_sleep_pool = True
class FakeAllocator:
def __init__(self):
self.pointer_to_data = {1: SimpleNamespace(handle=(0, 1024, 1, []))}
self.allocator_and_pools = {"weights": object()}
self.release_calls = 0
def get_current_usage(self):
return sum(data.handle[1] for data in self.pointer_to_data.values())
def release_pools(self):
assert worker.model_runner is None
assert worker.lora_manager is None
assert worker.profiler is None
assert worker._sleep_saved_buffers == {}
assert worker._step_lora_state == {}
assert self.get_current_usage() == 0
self.allocator_and_pools.clear()
self.release_calls += 1
events.append("pools")
allocator = FakeAllocator()
mocker.patch(
"vllm_omni.diffusion.worker.diffusion_worker._get_cumem_allocator_class",
return_value=SimpleNamespace(get_instance=lambda: allocator),
)
mocker.patch(
"vllm_omni.diffusion.worker.diffusion_worker.shutdown_kv_connector",
side_effect=lambda: events.append("connector"),
)
mocker.patch(
"vllm_omni.diffusion.worker.diffusion_worker.destroy_distributed_env",
side_effect=lambda: events.append("distributed"),
)
synchronize = mocker.patch(
"vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform.synchronize",
)
device_synchronized = False
def synchronize_device():
nonlocal device_synchronized
device_synchronized = True
events.append("synchronize")
synchronize.side_effect = synchronize_device
def collect_garbage():
events.append("gc")
if device_synchronized and worker.model_runner is None:
allocator.pointer_to_data.clear()
collect = mocker.patch(
"vllm_omni.diffusion.worker.diffusion_worker.gc.collect",
side_effect=collect_garbage,
)
worker.shutdown()
worker.shutdown()
assert events == ["offload", "kv", "distributed"]
assert events == [
"offload",
"kv",
"connector",
"distributed",
"synchronize",
"gc",
"pools",
"gc",
]
assert allocator.release_calls == 1
assert allocator.get_current_usage() == 0
assert allocator.allocator_and_pools == {}
assert not worker._owns_sleep_pool
kv_manager.close.assert_called_once_with()
destroy.assert_called_once_with()
synchronize.assert_called_once_with()
assert collect.call_count == 2
@pytest.mark.parametrize("cleanup_fails", [False, True])
def test_worker_shutdown_releases_model_references(self, mocker: MockerFixture, cleanup_fails):
@@ -272,11 +401,7 @@ class TestWorkerWrapperBaseDelegation:
worker._sleep_saved_buffers = {"buffer": weights}
del weights, pipeline
if cleanup_fails:
with pytest.raises(RuntimeError, match="offload cleanup failed"):
worker.shutdown()
else:
worker.shutdown()
worker.shutdown()
gc.collect()
assert worker.model_runner is None
@@ -297,11 +422,7 @@ class TestWorkerWrapperBaseDelegation:
worker.model_runner = None
worker._owns_sleep_pool = owns_pool
if owns_pool and cleanup_fails:
with pytest.raises(RuntimeError, match="pool release failed"):
worker.shutdown()
else:
worker.shutdown()
worker.shutdown()
assert release.call_count == int(owns_pool)
shutdown_connector.assert_called_once_with()
vllm_omni/diffusion/worker/diffusion_worker.py
+140 -80
View File
@@ -267,44 +267,52 @@ class DiffusionWorker:
# request id. Used by step mode to recover LoRA identity for cached
# requests, which only carry their request_id in subsequent ticks.
self._step_lora_state: dict[str, tuple[LoRARequest | None, float]] = {}
self._shutdown_complete = False
self.stage_id = getattr(od_config, "stage_id", 0)
self.init_device()
# Create model runner — one decision chain, in precedence order:
# 1. explicit od_config.diffusion_model_runner_cls (user override),
# 2. the runner declared by the engine class that engine_backend
# selects (e.g. ARDiffusionEngine -> ARDiffusionModelRunner),
# 3. the platform default.
# Routing policy therefore lives on the engine class / config surface;
# engines never mutate od_config. Overrides must be import-path
# strings — guard with isinstance so a non-string (e.g. a Mock
# od_config in tests) doesn't shadow the platform hook.
runner_override = getattr(self.od_config, "diffusion_model_runner_cls", None)
engine_runner = None
if not (isinstance(runner_override, str) and runner_override):
try:
from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
try:
# Create model runner — one decision chain, in precedence order:
# 1. explicit od_config.diffusion_model_runner_cls (user override),
# 2. the runner declared by the engine class that engine_backend
# selects (e.g. ARDiffusionEngine -> ARDiffusionModelRunner),
# 3. the platform default.
# Routing policy therefore lives on the engine class / config surface;
# engines never mutate od_config. Overrides must be import-path
# strings — guard with isinstance so a non-string (e.g. a Mock
# od_config in tests) doesn't shadow the platform hook.
runner_override = getattr(self.od_config, "diffusion_model_runner_cls", None)
engine_runner = None
if not (isinstance(runner_override, str) and runner_override):
try:
from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
engine_cls = DiffusionEngine.resolve_engine_class(self.od_config)
engine_runner = getattr(engine_cls, "default_diffusion_model_runner_cls", None)
except Exception:
logger.warning("Worker %s: engine_backend resolution failed; using platform runner", self.rank)
engine_runner = None
if isinstance(runner_override, str) and runner_override:
model_runner_cls_path = runner_override
elif isinstance(engine_runner, str) and engine_runner:
model_runner_cls_path = engine_runner
else:
model_runner_cls_path = current_omni_platform.get_diffusion_model_runner_cls()
model_runner_cls = resolve_obj_by_qualname(model_runner_cls_path)
self.model_runner = model_runner_cls(
vllm_config=self.vllm_config,
od_config=self.od_config,
device=self.device,
)
self.profiler: WorkerProfiler | None = self._create_profiler()
if not skip_load_model:
self.load_model(load_format=self.od_config.diffusion_load_format)
self.init_lora_manager()
engine_cls = DiffusionEngine.resolve_engine_class(self.od_config)
engine_runner = getattr(engine_cls, "default_diffusion_model_runner_cls", None)
except Exception:
logger.warning("Worker %s: engine_backend resolution failed; using platform runner", self.rank)
engine_runner = None
if isinstance(runner_override, str) and runner_override:
model_runner_cls_path = runner_override
elif isinstance(engine_runner, str) and engine_runner:
model_runner_cls_path = engine_runner
else:
model_runner_cls_path = current_omni_platform.get_diffusion_model_runner_cls()
model_runner_cls = resolve_obj_by_qualname(model_runner_cls_path)
self.model_runner = model_runner_cls(
vllm_config=self.vllm_config,
od_config=self.od_config,
device=self.device,
)
self.profiler: WorkerProfiler | None = self._create_profiler()
if not skip_load_model:
self.load_model(load_format=self.od_config.diffusion_load_format)
self.init_lora_manager()
except Exception:
# init_device() owns process-global distributed state. If anything
# after it fails, unwind that state before another inline worker is
# allowed to initialize in this process.
self.shutdown()
raise
logger.info(f"Worker {self.rank}: Initialization complete.")
def init_device(self) -> None:
@@ -1017,40 +1025,89 @@ class DiffusionWorker:
return nullcontext()
def shutdown(self) -> None:
"""Shutdown the worker and cleanup distributed environment."""
try:
if self.model_runner is not None:
mgr = getattr(self.model_runner, "kv_transfer_manager", None)
if mgr is None:
mgr = getattr(self.model_runner, "_kv_transfer_manager", None)
try:
offload_backend = getattr(self.model_runner, "offload_backend", None)
if offload_backend is not None:
offload_backend.disable()
finally:
if mgr is not None:
mgr.close()
finally:
self.model_runner = None
self.lora_manager = None
self._sleep_saved_buffers = {}
"""Shutdown the worker and release process-global resources."""
if getattr(self, "_shutdown_complete", False):
return
self._shutdown_complete = True
# Detach every model-owned reference before releasing the CuMem pools.
# Inline executors stay in the parent process, so relying on their
# wrapper becoming unreachable is not sufficient for deterministic
# teardown between sequential engine instances.
model_runner = getattr(self, "model_runner", None)
self.model_runner = None
self.lora_manager = None
self.profiler = None
self._sleep_saved_buffers = {}
self._step_lora_state = {}
if model_runner is not None:
mgr = getattr(model_runner, "kv_transfer_manager", None)
if mgr is None:
mgr = getattr(model_runner, "_kv_transfer_manager", None)
try:
if getattr(self, "_owns_sleep_pool", False):
offload_backend = getattr(model_runner, "offload_backend", None)
if offload_backend is not None:
offload_backend.disable()
except Exception:
logger.exception("Failed to disable diffusion offload backend during shutdown")
try:
if mgr is not None:
mgr.close()
except Exception:
logger.exception("Failed to close diffusion KV transfer manager during shutdown")
try:
shutdown_kv_connector()
except Exception:
logger.exception("Failed to shutdown diffusion KV connector")
try:
a2a_permute = sys.modules.get("vllm_omni.diffusion.distributed.a2a_permute")
if a2a_permute is not None:
a2a_permute.clear_a2a_permute_workspaces()
except Exception:
logger.exception("Failed to release fused Ulysses symmetric-memory workspaces")
try:
destroy_distributed_env()
except Exception:
logger.exception("Failed to destroy diffusion distributed environment")
# Drop the local runner only after its background services are stopped.
# Tensor finalizers must run before release_pools(), otherwise the
# singleton keeps allocations visible to the next sleep-enabled worker.
del model_runner
owns_sleep_pool = getattr(self, "_owns_sleep_pool", False)
if owns_sleep_pool:
try:
current_omni_platform.synchronize()
except Exception:
logger.exception("Failed to synchronize diffusion device during shutdown")
try:
gc.collect()
except Exception:
logger.exception("Failed to collect diffusion model resources during shutdown")
if owns_sleep_pool:
try:
CuMemAllocator = _get_cumem_allocator_class()
allocator = CuMemAllocator.get_instance()
if allocator.get_current_usage():
gc.collect()
_get_cumem_allocator_class().get_instance().release_pools()
self._owns_sleep_pool = False
finally:
try:
shutdown_kv_connector()
finally:
try:
a2a_permute = sys.modules.get("vllm_omni.diffusion.distributed.a2a_permute")
if a2a_permute is not None:
a2a_permute.clear_a2a_permute_workspaces()
except Exception:
logger.exception("Failed to release fused Ulysses symmetric-memory workspaces")
finally:
destroy_distributed_env()
allocator.release_pools()
self._owns_sleep_pool = False
# Re-collect after dropping the pool references so their
# finalizers remove every tracked allocation before the next
# inline worker checks get_current_usage().
gc.collect()
remaining_usage = allocator.get_current_usage()
if remaining_usage:
logger.warning(
"Diffusion CuMem allocator still tracks %s bytes after shutdown",
remaining_usage,
)
except Exception:
logger.exception("Failed to release diffusion CuMem pools during shutdown")
class CustomPipelineWorkerExtension:
@@ -1643,26 +1700,29 @@ class WorkerWrapperBase:
# Prepare worker class with extension support
worker_class = self._prepare_worker_class()
# Create the actual worker instance
# Only dynamic custom pipelines skip initial loading; native pipelines
# may also use custom_pipeline_args for model-specific component paths.
self.worker = worker_class(
# Create the actual worker instance. Only dynamic custom pipelines skip
# initial loading; native pipelines may also use custom_pipeline_args
# for model-specific component paths.
worker = worker_class(
local_rank=gpu_id,
rank=gpu_id,
od_config=od_config,
skip_load_model=self.uses_custom_pipeline,
)
# Re-initialize pipeline with custom pipeline if provided
if self.uses_custom_pipeline:
try:
# Re-initialize pipeline with custom pipeline if provided.
if self.uses_custom_pipeline:
worker.re_init_pipeline(self.custom_pipeline_args)
except Exception:
# The executor never receives a wrapper whose constructor raises.
# Unwind a successfully created worker here so its distributed
# groups do not poison the next inline engine initialization.
try:
self.worker.re_init_pipeline(self.custom_pipeline_args)
worker.shutdown()
except Exception:
try:
self.worker.shutdown()
except Exception:
logger.exception("Failed to clean up worker after custom pipeline initialization failure")
raise
logger.exception("Failed to clean up worker after custom pipeline initialization failure")
raise
self.worker = worker
def _prepare_worker_class(self) -> type:
"""