add check_health in inline client (#3052)
Signed-off-by: rongfu.leng <lenronfu@gmail.com>
This commit is contained in:
1 parent
0309b3bd19
commit
da60663f78
2 files changed
+8
-1
No files matched your search
@@ -385,9 +385,9 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
|
||||
raise
|
||||
|
||||
def check_health(self) -> None:
|
||||
self._ensure_open()
|
||||
if self.is_failed:
|
||||
raise EngineDeadError()
|
||||
self._ensure_open()
|
||||
for p in self._processes:
|
||||
if not p.is_alive():
|
||||
self.is_failed = True
|
||||
|
||||
@@ -15,6 +15,7 @@ from typing import TYPE_CHECKING, Any
|
||||
import torch
|
||||
from PIL import Image
|
||||
from vllm.logger import init_logger
|
||||
from vllm.v1.engine.exceptions import EngineDeadError
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionRequestAbortedError
|
||||
from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
|
||||
@@ -329,6 +330,12 @@ class InlineStageDiffusionClient:
|
||||
None,
|
||||
)
|
||||
|
||||
def check_health(self) -> None:
|
||||
"""Check if the inline diffusion engine and its workers are healthy."""
|
||||
if self._shutting_down:
|
||||
raise EngineDeadError("InlineStageDiffusionClient is shutting down")
|
||||
self._engine.executor.check_health()
|
||||
|
||||
def shutdown(self) -> None:
|
||||
self._shutting_down = True
|
||||
|
||||
|
||||
Reference in new issue
Block a user