[Bugfix][Diffusion] Clean up inline worker lifecycle (#7986)
Signed-off-by: andyluo7 <andy.luo@amd.com>
This commit is contained in:
andyluo7 authored and GitHub committed
1 parent
0a0e1884f7
commit
9b08ac8203
2 files changed
+278
-97
No files matched your search
@@ -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:
|
||||
"""
|
||||
|
||||
Reference in new issue
Block a user