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

[release/v0.18.0.post1][Bugfix] fix oom no response to user (#2638)

Browse Source 
Signed-off-by: bjf-frz <frz123db@gmail.com>
This commit is contained in:
bjf-frz authored and GitHub committed 2026-04-09 20:31:58 +08:00
1 parent 589e664947
commit ca7ee8d6e7
12 files changed
+332 -47
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

vllm_omni/diffusion/diffusion_engine.py
+50 -26
View File
@@ -11,7 +11,7 @@ import PIL.Image
import torch
from vllm.logger import init_logger
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig, OmniRequestError, normalize_omni_error
from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
from vllm_omni.diffusion.registry import (
DiffusionModelRegistry,
@@ -99,7 +99,12 @@ class DiffusionEngine:
exec_total_time = time.perf_counter() - exec_start_time
if output.error:
raise Exception(f"{output.error}")
# raise Exception(f"{output.error}")
raise OmniRequestError(
output.error,
status_code=500,
error_type="DiffusionExecutionError",
)
logger.info("Generation completed successfully.")
if output.output is None:
@@ -280,33 +285,52 @@ class DiffusionEngine:
target_sched_req_id = self.scheduler.add_request(request)
# keep scheduling and executing until the target request is finished
while True:
sched_output = self.scheduler.schedule()
if sched_output.is_empty:
if not self.scheduler.has_requests():
raise RuntimeError("Diffusion scheduler has no runnable requests.")
continue
try:
while True:
sched_output = self.scheduler.schedule()
if sched_output.is_empty:
if not self.scheduler.has_requests():
# raise RuntimeError("Diffusion scheduler has no runnable requests.")
raise OmniRequestError(
"Diffusion scheduler has no runnable requests.",
status_code=500,
error_type="SchedulerError",
)
continue
# NOTE: add_req_and_wait_for_response() is synchronous, and
# the scheduler currently enforces _max_batch_size = 1 (see
# vllm_omni/diffusion/sched/base_scheduler.py), so we directly
# take the single scheduled request here.
sched_req_id = sched_output.scheduled_req_ids[0]
req = sched_output.scheduled_new_reqs[0].req
# NOTE: add_req_and_wait_for_response() is synchronous, and
# the scheduler currently enforces _max_batch_size = 1 (see
# vllm_omni/diffusion/sched/base_scheduler.py), so we directly
# take the single scheduled request here.
sched_req_id = sched_output.scheduled_req_ids[0]
req = sched_output.scheduled_new_reqs[0].req
try:
output = self.executor.add_req(req)
except Exception as exc:
logger.error(
"Execution failed for diffusion request %s",
sched_req_id,
exc_info=True,
)
# output = DiffusionOutput(error=str(exc))
raise normalize_omni_error(exc) from exc
finished_req_ids = self.scheduler.update_from_output(sched_output, output)
if output.error:
raise OmniRequestError(
output.error,
status_code=500,
error_type="DiffusionExecutionError",
)
if target_sched_req_id in finished_req_ids:
# self.scheduler.pop_request_state(target_sched_req_id)
return output
finally:
try:
output = self.executor.add_req(req)
except Exception as exc:
logger.error(
"Execution failed for diffusion request %s",
sched_req_id,
exc_info=True,
)
output = DiffusionOutput(error=str(exc))
finished_req_ids = self.scheduler.update_from_output(sched_output, output)
if target_sched_req_id in finished_req_ids:
self.scheduler.pop_request_state(target_sched_req_id)
return output
except Exception:
logger.debug("Request state already removed: %s", target_sched_req_id, exc_info=True)
def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None:
"""Start or stop torch profiling on all diffusion workers.