693 lines
30 KiB
Python
693 lines
30 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import inspect
|
|
import queue
|
|
import threading
|
|
import time
|
|
from collections.abc import Iterable
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import PIL.Image
|
|
import torch
|
|
from vllm.logger import init_logger
|
|
from vllm.v1.engine.exceptions import EngineDeadError
|
|
|
|
from vllm_omni.diffusion.data import (
|
|
DiffusionOutput,
|
|
DiffusionRequestAbortedError,
|
|
OmniDiffusionConfig,
|
|
)
|
|
from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
|
|
from vllm_omni.diffusion.registry import (
|
|
DiffusionModelRegistry,
|
|
get_diffusion_post_process_func,
|
|
get_diffusion_pre_process_func,
|
|
)
|
|
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
|
from vllm_omni.diffusion.sched import RequestScheduler, SchedulerInterface, StepScheduler
|
|
from vllm_omni.diffusion.sched.interface import DiffusionRequestStatus
|
|
from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt
|
|
from vllm_omni.outputs import OmniRequestOutput
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
def supports_multimodal_input(od_config: OmniDiffusionConfig) -> tuple[bool, bool]:
|
|
if od_config.diffusion_load_format == "diffusers" and (pipe_cls := od_config.diffusers_pipeline_cls) is not None:
|
|
signature = inspect.signature(pipe_cls.__call__)
|
|
support_image_input = "image" in signature.parameters
|
|
support_audio_input = (
|
|
"audio" in signature.parameters or "audio_latents" in signature.parameters
|
|
) # ref. LTX-2 format
|
|
return support_image_input, support_audio_input
|
|
|
|
supports_image_input = False
|
|
supports_audio_input = False
|
|
|
|
model_cls = DiffusionModelRegistry._try_load_model_cls(od_config.model_class_name)
|
|
if model_cls is not None:
|
|
supports_image_input = bool(getattr(model_cls, "support_image_input", False))
|
|
supports_audio_input = bool(getattr(model_cls, "support_audio_input", False))
|
|
return supports_image_input, supports_audio_input
|
|
|
|
|
|
def image_color_format(model_class_name: str) -> str:
|
|
model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)
|
|
return getattr(model_cls, "color_format", "RGB")
|
|
|
|
|
|
def supports_audio_output(model_class_name: str) -> bool:
|
|
model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)
|
|
if model_cls is None:
|
|
return False
|
|
return bool(getattr(model_cls, "support_audio_output", False))
|
|
|
|
|
|
class DiffusionEngine:
|
|
"""The diffusion engine for vLLM-Omni diffusion models."""
|
|
|
|
def __init__(
|
|
self,
|
|
od_config: OmniDiffusionConfig,
|
|
scheduler: SchedulerInterface | None = None,
|
|
):
|
|
"""Initialize the diffusion engine.
|
|
|
|
Args:
|
|
config: The configuration for the diffusion engine.
|
|
"""
|
|
self.od_config = od_config
|
|
|
|
self.post_process_func = get_diffusion_post_process_func(od_config)
|
|
self.pre_process_func = get_diffusion_pre_process_func(od_config)
|
|
# Cache whether the model-specific postprocess accepts request-level
|
|
# sampling params so step() can support both legacy and extended hooks.
|
|
self._post_process_accepts_sampling_params = bool(
|
|
self.post_process_func is not None
|
|
and "sampling_params" in inspect.signature(self.post_process_func).parameters
|
|
)
|
|
|
|
executor_class = DiffusionExecutor.get_class(od_config)
|
|
self.executor = executor_class(od_config)
|
|
self.step_execution = bool(getattr(od_config, "step_execution", False))
|
|
self.scheduler: SchedulerInterface = scheduler or (
|
|
StepScheduler() if self.step_execution else RequestScheduler()
|
|
)
|
|
self.scheduler.initialize(od_config)
|
|
if self.scheduler.max_num_running_reqs > 1 and not self.step_execution:
|
|
max_num_seqs = self.scheduler.max_num_running_reqs
|
|
self.scheduler.max_num_running_reqs = 1
|
|
logger.warning(f"Non-stepwise-execution does not support max-num-seqs={max_num_seqs}, set it to 1.")
|
|
self.main_loop: asyncio.AbstractEventLoop | None = None
|
|
self.stop_event: threading.Event | None = None
|
|
self._loop_started = False
|
|
self._init_lock = asyncio.Lock()
|
|
self._rpc_lock = threading.RLock()
|
|
self._cv = threading.Condition(self._rpc_lock)
|
|
self._out_queue: dict[str, asyncio.Future] = {}
|
|
self.abort_queue: queue.Queue[str] = queue.Queue()
|
|
self.execute_fn = self.executor.execute_step if self.step_execution else self.executor.execute_request
|
|
|
|
try:
|
|
self._dummy_run()
|
|
except Exception as e:
|
|
logger.error(f"Dummy run failed: {e}")
|
|
self.close()
|
|
raise e
|
|
|
|
async def _check_and_start_background_loop(self):
|
|
if self._loop_started:
|
|
return
|
|
|
|
async with self._init_lock:
|
|
# double check, in case of lock queue issue
|
|
if self._loop_started:
|
|
return
|
|
|
|
self.main_loop = asyncio.get_running_loop()
|
|
self.stop_event = threading.Event()
|
|
self.worker_thread = threading.Thread(target=self._busy_loop)
|
|
self.worker_thread.start()
|
|
self._loop_started = True
|
|
|
|
async def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]:
|
|
await self._check_and_start_background_loop()
|
|
|
|
diffusion_engine_start_time = time.perf_counter()
|
|
|
|
# Apply pre-processing if available
|
|
preprocess_time = 0.0
|
|
if self.pre_process_func is not None:
|
|
preprocess_start_time = time.perf_counter()
|
|
request = self.pre_process_func(request)
|
|
preprocess_time = time.perf_counter() - preprocess_start_time
|
|
logger.debug("Pre-processing completed in %.4f seconds", preprocess_time)
|
|
|
|
exec_start_time = time.perf_counter()
|
|
output = await self.async_add_req_and_wait_for_response(request)
|
|
exec_total_time = time.perf_counter() - exec_start_time
|
|
|
|
if output.aborted:
|
|
raise DiffusionRequestAbortedError(output.abort_message or "Diffusion request aborted.")
|
|
if output.error:
|
|
raise RuntimeError(output.error)
|
|
logger.debug("Generation completed successfully.")
|
|
|
|
if output.output is None:
|
|
logger.warning("Output is None, returning empty OmniRequestOutput")
|
|
return [
|
|
OmniRequestOutput.from_diffusion(
|
|
request_id=request.request_ids[i] if i < len(request.request_ids) else "",
|
|
images=[],
|
|
prompt=prompt,
|
|
metrics={},
|
|
latents=None,
|
|
)
|
|
for i, prompt in enumerate(request.prompts)
|
|
]
|
|
|
|
# When CPU offload is enabled, move output to CPU before
|
|
# post-processing to avoid device OOM — model weights may still
|
|
# reside on the device and leave no headroom for intermediates.
|
|
output_data = output.output
|
|
if (
|
|
self.od_config.enable_cpu_offload
|
|
and isinstance(output_data, torch.Tensor)
|
|
and output_data.device.type != "cpu"
|
|
):
|
|
output_data = output_data.cpu()
|
|
|
|
postprocess_start_time = time.perf_counter()
|
|
if self.post_process_func is not None:
|
|
# Some video pipelines need request-level controls during
|
|
# postprocess (for example worker-side frame interpolation).
|
|
if self._post_process_accepts_sampling_params:
|
|
outputs = self.post_process_func(output_data, sampling_params=request.sampling_params)
|
|
else:
|
|
outputs = self.post_process_func(output_data)
|
|
else:
|
|
outputs = output_data
|
|
audio_payload = None
|
|
custom_output = output.custom_output or {}
|
|
model_audio_sample_rate = None
|
|
model_fps = None
|
|
if isinstance(outputs, dict):
|
|
audio_payload = outputs.get("audio")
|
|
custom_output.update(outputs.get("custom_output") or {})
|
|
model_audio_sample_rate = outputs.get("audio_sample_rate")
|
|
model_fps = outputs.get("fps")
|
|
outputs = outputs.get("video", outputs)
|
|
postprocess_time = time.perf_counter() - postprocess_start_time
|
|
logger.debug("Post-processing completed in %.4f seconds", postprocess_time)
|
|
|
|
step_total_ms = (time.perf_counter() - diffusion_engine_start_time) * 1000
|
|
logger.debug(
|
|
"DiffusionEngine.step breakdown: preprocess=%.2f ms, "
|
|
"add_req_and_wait=%.2f ms, postprocess=%.2f ms, total=%.2f ms",
|
|
preprocess_time * 1000,
|
|
exec_total_time * 1000,
|
|
postprocess_time * 1000,
|
|
step_total_ms,
|
|
)
|
|
|
|
# Convert to OmniRequestOutput format
|
|
# Ensure outputs is a list
|
|
if not isinstance(outputs, list):
|
|
outputs = [outputs] if outputs is not None else []
|
|
|
|
metrics = {
|
|
"preprocess_time_ms": preprocess_time * 1000,
|
|
"diffusion_engine_exec_time_ms": exec_total_time * 1000,
|
|
"diffusion_engine_total_time_ms": step_total_ms,
|
|
"image_num": int(request.sampling_params.num_outputs_per_prompt),
|
|
"resolution": int(request.sampling_params.resolution),
|
|
"postprocess_time_ms": postprocess_time * 1000,
|
|
}
|
|
|
|
# Handle single request or multiple requests
|
|
is_audio_output = supports_audio_output(self.od_config.model_class_name)
|
|
if is_audio_output and model_audio_sample_rate is None:
|
|
model_cls = DiffusionModelRegistry._try_load_model_cls(self.od_config.model_class_name)
|
|
model_audio_sample_rate = getattr(model_cls, "audio_sample_rate", None)
|
|
|
|
def _audio_mm(payload: Any) -> dict[str, Any]:
|
|
mm: dict[str, Any] = {"audio": payload}
|
|
if model_audio_sample_rate is not None:
|
|
mm["audio_sample_rate"] = model_audio_sample_rate
|
|
return mm
|
|
|
|
if len(request.prompts) == 1:
|
|
# Single request: return single OmniRequestOutput
|
|
prompt = request.prompts[0]
|
|
request_id = request.request_ids[0] if request.request_ids else ""
|
|
|
|
if is_audio_output:
|
|
request_audio_payload = outputs[0] if len(outputs) == 1 else outputs
|
|
return [
|
|
OmniRequestOutput.from_diffusion(
|
|
request_id=request_id,
|
|
images=[],
|
|
prompt=prompt,
|
|
metrics=metrics,
|
|
latents=output.trajectory_latents,
|
|
trajectory_latents=output.trajectory_latents,
|
|
trajectory_timesteps=output.trajectory_timesteps,
|
|
trajectory_log_probs=output.trajectory_log_probs,
|
|
trajectory_decoded=output.trajectory_decoded,
|
|
multimodal_output=_audio_mm(request_audio_payload),
|
|
final_output_type="audio",
|
|
stage_durations=output.stage_durations,
|
|
peak_memory_mb=output.peak_memory_mb,
|
|
),
|
|
]
|
|
else:
|
|
mm_output = {}
|
|
if audio_payload is not None:
|
|
mm_output["audio"] = audio_payload
|
|
if model_audio_sample_rate is not None:
|
|
mm_output["audio_sample_rate"] = model_audio_sample_rate
|
|
if model_fps is not None:
|
|
mm_output["fps"] = model_fps
|
|
return [
|
|
OmniRequestOutput.from_diffusion(
|
|
request_id=request_id,
|
|
images=outputs,
|
|
prompt=prompt,
|
|
metrics=metrics,
|
|
latents=output.trajectory_latents,
|
|
trajectory_latents=output.trajectory_latents,
|
|
trajectory_timesteps=output.trajectory_timesteps,
|
|
trajectory_log_probs=output.trajectory_log_probs,
|
|
trajectory_decoded=output.trajectory_decoded,
|
|
custom_output=custom_output,
|
|
multimodal_output=mm_output,
|
|
stage_durations=output.stage_durations,
|
|
peak_memory_mb=output.peak_memory_mb,
|
|
),
|
|
]
|
|
else:
|
|
# Multiple requests: return list of OmniRequestOutput
|
|
# Split images based on num_outputs_per_prompt for each request
|
|
results = []
|
|
output_idx = 0
|
|
|
|
for i, prompt in enumerate(request.prompts):
|
|
request_id = request.request_ids[i] if i < len(request.request_ids) else ""
|
|
|
|
# Get images for this request
|
|
num_outputs = request.sampling_params.num_outputs_per_prompt
|
|
start_idx = output_idx
|
|
end_idx = start_idx + num_outputs
|
|
request_outputs = outputs[start_idx:end_idx] if output_idx < len(outputs) else []
|
|
output_idx = end_idx
|
|
|
|
if is_audio_output:
|
|
request_audio_payload = request_outputs[0] if len(request_outputs) == 1 else request_outputs
|
|
results.append(
|
|
OmniRequestOutput.from_diffusion(
|
|
request_id=request_id,
|
|
images=[],
|
|
prompt=prompt,
|
|
metrics=metrics,
|
|
latents=output.trajectory_latents,
|
|
trajectory_latents=output.trajectory_latents,
|
|
trajectory_timesteps=output.trajectory_timesteps,
|
|
trajectory_log_probs=output.trajectory_log_probs,
|
|
trajectory_decoded=output.trajectory_decoded,
|
|
multimodal_output=_audio_mm(request_audio_payload),
|
|
final_output_type="audio",
|
|
stage_durations=output.stage_durations,
|
|
peak_memory_mb=output.peak_memory_mb,
|
|
),
|
|
)
|
|
else:
|
|
mm_output = {}
|
|
if audio_payload is not None:
|
|
sliced_audio = audio_payload
|
|
if isinstance(audio_payload, (list, tuple)):
|
|
sliced_audio = audio_payload[start_idx:end_idx]
|
|
if len(sliced_audio) == 1:
|
|
sliced_audio = sliced_audio[0]
|
|
elif hasattr(audio_payload, "shape") and getattr(audio_payload, "shape", None) is not None:
|
|
if len(audio_payload.shape) > 0 and audio_payload.shape[0] >= end_idx:
|
|
sliced_audio = audio_payload[start_idx:end_idx]
|
|
if num_outputs == 1:
|
|
sliced_audio = sliced_audio[0]
|
|
mm_output["audio"] = sliced_audio
|
|
if model_audio_sample_rate is not None:
|
|
mm_output["audio_sample_rate"] = model_audio_sample_rate
|
|
if model_fps is not None:
|
|
mm_output["fps"] = model_fps
|
|
results.append(
|
|
OmniRequestOutput.from_diffusion(
|
|
request_id=request_id,
|
|
images=request_outputs,
|
|
prompt=prompt,
|
|
metrics=metrics,
|
|
latents=output.trajectory_latents,
|
|
trajectory_latents=output.trajectory_latents,
|
|
trajectory_timesteps=output.trajectory_timesteps,
|
|
trajectory_log_probs=output.trajectory_log_probs,
|
|
trajectory_decoded=output.trajectory_decoded,
|
|
custom_output=custom_output,
|
|
multimodal_output=mm_output,
|
|
stage_durations=output.stage_durations,
|
|
peak_memory_mb=output.peak_memory_mb,
|
|
),
|
|
)
|
|
|
|
return results
|
|
|
|
def _busy_loop(self):
|
|
while not self.stop_event.is_set():
|
|
self._process_aborts_queue()
|
|
|
|
with self._cv:
|
|
while not self.scheduler.has_requests() and not self.stop_event.is_set():
|
|
self._cv.wait(timeout=1.0)
|
|
|
|
if self.stop_event.is_set():
|
|
break
|
|
|
|
sched_output = self.scheduler.schedule()
|
|
|
|
if sched_output.is_empty:
|
|
self._handle_finished_requests(sched_output.finished_req_ids, None)
|
|
continue
|
|
|
|
try:
|
|
runner_output = self.execute_fn(sched_output)
|
|
except Exception as exc:
|
|
logger.error(
|
|
"Execution failed for diffusion requests %s", sched_output.scheduled_req_ids, exc_info=True
|
|
)
|
|
runner_output = BatchRunnerOutput.from_list(
|
|
[
|
|
RunnerOutput(
|
|
req_id=req_id,
|
|
step_index=None,
|
|
finished=True,
|
|
result=DiffusionOutput(error=str(exc)),
|
|
)
|
|
for req_id in sched_output.scheduled_req_ids
|
|
]
|
|
)
|
|
|
|
self._process_aborts_queue()
|
|
finished_req_ids = self.scheduler.update_from_output(sched_output, runner_output)
|
|
self._handle_finished_requests(finished_req_ids, runner_output)
|
|
|
|
def _handle_finished_requests(
|
|
self,
|
|
finished_ids: set[str],
|
|
runner_output: RunnerOutput | None = None,
|
|
missing_result_error: str = "Diffusion execution finished without a final output",
|
|
):
|
|
for rid in finished_ids:
|
|
if rid in self._out_queue:
|
|
fut = self._out_queue.pop(rid)
|
|
if runner_output is not None:
|
|
_output = runner_output.get_req_output(rid)
|
|
else:
|
|
_output = None
|
|
out = self._finalize_finished_request(rid, _output, missing_result_error)
|
|
if not fut.done():
|
|
self.main_loop.call_soon_threadsafe(fut.set_result, out)
|
|
|
|
@staticmethod
|
|
def make_engine(
|
|
config: OmniDiffusionConfig,
|
|
scheduler: SchedulerInterface | None = None,
|
|
) -> DiffusionEngine:
|
|
"""Factory method to create a DiffusionEngine instance.
|
|
|
|
Args:
|
|
config: The configuration for the diffusion engine.
|
|
|
|
Returns:
|
|
An instance of DiffusionEngine.
|
|
"""
|
|
return DiffusionEngine(config, scheduler=scheduler)
|
|
|
|
def add_request(self, request: OmniDiffusionRequest) -> str:
|
|
fut = self.main_loop.create_future()
|
|
with self._cv:
|
|
sched_req_id = self.scheduler.add_request(request)
|
|
self._out_queue[sched_req_id] = fut
|
|
self._cv.notify_all()
|
|
|
|
return sched_req_id
|
|
|
|
async def get_result(self, request_id: str) -> DiffusionOutput:
|
|
fut = self._out_queue.get(request_id)
|
|
|
|
if fut is None:
|
|
raise RuntimeError(f"Request {request_id} not found in output queue.")
|
|
try:
|
|
return await fut
|
|
except Exception as e:
|
|
logger.error(f"Wait for response failed: {e}")
|
|
raise
|
|
|
|
async def async_add_req_and_wait_for_response(self, request: OmniDiffusionRequest) -> DiffusionOutput:
|
|
sched_req_id = self.add_request(request)
|
|
return await self.get_result(sched_req_id)
|
|
|
|
def add_req_and_wait_for_response(self, request: OmniDiffusionRequest) -> DiffusionOutput:
|
|
with self._rpc_lock:
|
|
target_sched_req_id = self.scheduler.add_request(request)
|
|
|
|
# keep scheduling and executing until the target request is finished
|
|
while True:
|
|
self._process_aborts_queue()
|
|
sched_output = self.scheduler.schedule()
|
|
if sched_output.is_empty:
|
|
if target_sched_req_id in sched_output.finished_req_ids:
|
|
return self._finalize_finished_request(target_sched_req_id)
|
|
if not self.scheduler.has_requests():
|
|
raise RuntimeError("Diffusion scheduler has no runnable requests.")
|
|
continue
|
|
|
|
# NOTE: add_req_and_wait_for_response() is synchronous, will be only called
|
|
# within _dummy_run, only one request will be scheduled
|
|
sched_req_id = sched_output.scheduled_req_ids[0]
|
|
try:
|
|
runner_output = self.execute_fn(sched_output)
|
|
except EngineDeadError:
|
|
raise
|
|
except Exception as exc:
|
|
logger.error("Execution failed for diffusion request %s", sched_req_id, exc_info=True)
|
|
runner_output = RunnerOutput(
|
|
req_id=sched_req_id,
|
|
step_index=None,
|
|
finished=True,
|
|
result=DiffusionOutput(error=str(exc)),
|
|
)
|
|
|
|
self._process_aborts_queue()
|
|
|
|
finished_req_ids = self.scheduler.update_from_output(sched_output, runner_output)
|
|
|
|
# sync func should receive one result
|
|
if not isinstance(runner_output, RunnerOutput) and not len(runner_output) == 1:
|
|
raise ValueError("Sync func should receive one result at one time")
|
|
if target_sched_req_id in finished_req_ids:
|
|
runner_output = runner_output.get_req_output(target_sched_req_id)
|
|
return self._finalize_finished_request(
|
|
target_sched_req_id,
|
|
runner_output=runner_output,
|
|
missing_result_error="Diffusion execution finished without a final output.",
|
|
)
|
|
|
|
def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None:
|
|
"""Start or stop profiling on all diffusion workers.
|
|
|
|
Args:
|
|
is_start: True to start profiling, False to stop.
|
|
profile_prefix: Optional prefix for trace filename.
|
|
"""
|
|
if is_start:
|
|
if profile_prefix is None:
|
|
profile_prefix = f"diffusion_{int(time.time())}"
|
|
logger.info(f"Starting diffusion profiling with prefix: {profile_prefix}")
|
|
else:
|
|
logger.info("Stopping diffusion profiling...")
|
|
|
|
try:
|
|
self.collective_rpc(method="profile", args=(is_start, profile_prefix))
|
|
except Exception as e:
|
|
action = "start" if is_start else "stop"
|
|
logger.error(f"Failed to {action} profiling on workers", exc_info=True)
|
|
if is_start:
|
|
raise RuntimeError(f"Could not {action} profiler: {e}") from e
|
|
|
|
def _dummy_run(self):
|
|
"""A dummy run to warm up the model."""
|
|
num_inference_steps = 1
|
|
height = 512
|
|
width = 512
|
|
prompt: OmniTextPrompt = {"prompt": "dummy run"}
|
|
|
|
supports_image_input, supports_audio_input = supports_multimodal_input(self.od_config)
|
|
if supports_image_input:
|
|
# Provide a dummy image input if the model supports it
|
|
color_format = image_color_format(self.od_config.model_class_name)
|
|
dummy_image = PIL.Image.new(color_format, (width, height))
|
|
prompt.setdefault("multi_modal_data", {})["image"] = dummy_image
|
|
|
|
if supports_audio_input:
|
|
audio_sr = 16000
|
|
dummy_audio = np.random.randn(audio_sr * 2).astype(np.float32)
|
|
prompt.setdefault("multi_modal_data", {})["audio"] = dummy_audio
|
|
|
|
req = OmniDiffusionRequest(
|
|
prompts=[prompt],
|
|
request_ids=["dummy_req_id"],
|
|
sampling_params=OmniDiffusionSamplingParams(
|
|
height=height,
|
|
width=width,
|
|
num_inference_steps=num_inference_steps,
|
|
# Keep warmup path minimal and robust across text encoders.
|
|
# Some models may fail when warmup implicitly triggers
|
|
# classifier-free guidance with an empty negative prompt.
|
|
guidance_scale=0.0,
|
|
num_outputs_per_prompt=1,
|
|
# Disable CFG for warmup to avoid triggering CFG parallel
|
|
# validation when cfg_parallel_size > 1.
|
|
extra_args={"cfg_text_scale": 1.0, "cfg_img_scale": 1.0},
|
|
),
|
|
)
|
|
logger.info("dummy run to warm up the model")
|
|
request = self.pre_process_func(req) if self.pre_process_func is not None else req
|
|
output = self.add_req_and_wait_for_response(request)
|
|
if output.error:
|
|
raise RuntimeError(f"Dummy run failed: {output.error}")
|
|
|
|
def collective_rpc(
|
|
self,
|
|
method: str,
|
|
timeout: float | None = None,
|
|
args: tuple = (),
|
|
kwargs: dict | None = None,
|
|
unique_reply_rank: int | None = None,
|
|
) -> Any:
|
|
"""Call a method on worker processes and get results immediately.
|
|
|
|
Args:
|
|
method: The method name (str) to execute on workers
|
|
timeout: Optional timeout in seconds
|
|
args: Positional arguments for the method
|
|
kwargs: Keyword arguments for the method
|
|
unique_reply_rank: If set, only get reply from this rank
|
|
|
|
Returns:
|
|
Single result if unique_reply_rank is provided, otherwise list of results
|
|
"""
|
|
assert isinstance(method, str), "Only string method names are supported for now"
|
|
|
|
deadline = None if timeout is None else time.monotonic() + timeout
|
|
acquired = False
|
|
try:
|
|
if deadline is None:
|
|
self._rpc_lock.acquire()
|
|
acquired = True
|
|
else:
|
|
lock_timeout = max(0, deadline - time.monotonic())
|
|
acquired = self._rpc_lock.acquire(timeout=lock_timeout)
|
|
if not acquired:
|
|
raise TimeoutError(f"RPC call to {method} timed out waiting for engine lock.")
|
|
|
|
rpc_timeout = None if deadline is None else max(0, deadline - time.monotonic())
|
|
if deadline is not None and rpc_timeout <= 0:
|
|
raise TimeoutError(f"RPC call to {method} timed out.")
|
|
|
|
return self.executor.collective_rpc(
|
|
method=method,
|
|
timeout=rpc_timeout,
|
|
args=args,
|
|
kwargs=kwargs,
|
|
unique_reply_rank=unique_reply_rank,
|
|
)
|
|
finally:
|
|
if acquired:
|
|
self._rpc_lock.release()
|
|
|
|
def close(self) -> None:
|
|
if hasattr(self, "scheduler"):
|
|
self.scheduler.close()
|
|
if hasattr(self, "stop_event") and self.stop_event is not None:
|
|
with self._cv:
|
|
self.stop_event.set()
|
|
self._cv.notify_all()
|
|
if hasattr(self, "worker_thread") and self.worker_thread.is_alive():
|
|
self.worker_thread.join(timeout=10)
|
|
if self.worker_thread.is_alive():
|
|
logger.warning("Worker thread did not terminate within 10s, resources may not be released.")
|
|
else:
|
|
self._loop_started = False
|
|
if hasattr(self, "executor"):
|
|
self.executor.shutdown()
|
|
|
|
def abort(self, request_id: str | Iterable[str]) -> None:
|
|
request_ids = [request_id] if isinstance(request_id, str) else list(request_id)
|
|
|
|
with self._cv:
|
|
for req_id in request_ids:
|
|
self.abort_queue.put(req_id)
|
|
self._cv.notify_all()
|
|
|
|
def _process_aborts_queue(self) -> None:
|
|
if self.abort_queue.empty():
|
|
return
|
|
|
|
request_ids: list[str] = []
|
|
while not self.abort_queue.empty():
|
|
ids = self.abort_queue.get_nowait()
|
|
request_ids.extend((ids,) if isinstance(ids, str) else ids)
|
|
|
|
self._abort_requests(request_ids)
|
|
|
|
def _abort_requests(self, request_ids: str | Iterable[str]) -> None:
|
|
request_ids = [request_ids] if isinstance(request_ids, str) else list(request_ids)
|
|
|
|
sched_req_ids: list[str] = []
|
|
for request_id in dict.fromkeys(request_ids):
|
|
sched_req_id = self.scheduler.get_sched_req_id(request_id)
|
|
if sched_req_id is not None:
|
|
sched_req_ids.append(sched_req_id)
|
|
|
|
for sched_req_id in dict.fromkeys(sched_req_ids):
|
|
if self.scheduler.get_request_state(sched_req_id) is not None:
|
|
self.scheduler.finish_requests(sched_req_id, DiffusionRequestStatus.FINISHED_ABORTED)
|
|
|
|
def _finalize_finished_request(
|
|
self,
|
|
sched_req_id: str,
|
|
runner_output: RunnerOutput | None = None,
|
|
missing_result_error: str = "Diffusion scheduler finished target request without execution output.",
|
|
) -> DiffusionOutput:
|
|
state = self.scheduler.get_request_state(sched_req_id)
|
|
popped_state = self.scheduler.pop_request_state(sched_req_id)
|
|
state = state or popped_state
|
|
|
|
if state is None:
|
|
raise RuntimeError(f"Diffusion scheduler lost state for request {sched_req_id}.")
|
|
|
|
if state.status == DiffusionRequestStatus.FINISHED_ABORTED:
|
|
request_id = state.req.request_ids[0] if state.req.request_ids else sched_req_id
|
|
return DiffusionOutput(
|
|
aborted=True,
|
|
abort_message=f"Request {request_id} aborted.",
|
|
)
|
|
|
|
if runner_output is not None and runner_output.result is not None:
|
|
return runner_output.result
|
|
|
|
return DiffusionOutput(error=missing_result_error)
|