560 lines
24 KiB
Python
560 lines
24 KiB
Python
import multiprocessing as mp
|
|
import os
|
|
import sys
|
|
import time
|
|
from collections.abc import Sequence
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
from typing import Any, Optional, Union
|
|
|
|
import cloudpickle
|
|
from pydantic import ValidationError
|
|
|
|
# External library imports (vLLM)
|
|
from vllm.config import CompilationConfig, StructuredOutputsConfig, is_init_field
|
|
from vllm.engine.arg_utils import HfOverrides
|
|
from vllm.entrypoints.llm import LLM
|
|
from vllm.inputs import PromptType
|
|
from vllm.logger import init_logger
|
|
from vllm.plugins.io_processors import get_io_processor
|
|
from vllm.sampling_params import SamplingParams
|
|
from vllm.usage.usage_lib import UsageContext
|
|
from vllm.utils import Counter
|
|
from vllm.v1.engine.llm_engine import LLMEngine
|
|
|
|
# Internal imports (our code)
|
|
from vllm_omni.engine.arg_utils import OmniEngineArgs
|
|
from vllm_omni.engine.output_processor import MultimodalOutputProcessor
|
|
from vllm_omni.engine.processor import OmniProcessor
|
|
from vllm_omni.entrypoints.log_utils import (
|
|
OrchestratorMetrics,
|
|
configure_orchestrator_logger,
|
|
init_stats_paths,
|
|
remove_old_logs,
|
|
)
|
|
from vllm_omni.entrypoints.omni_stage import OmniStage
|
|
from vllm_omni.entrypoints.stage_utils import encode_for_ipc as _encode
|
|
from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load
|
|
from vllm_omni.entrypoints.stage_utils import serialize_obj as _set
|
|
from vllm_omni.entrypoints.utils import load_stage_configs_from_model, load_stage_configs_from_yaml
|
|
from vllm_omni.outputs import OmniRequestOutput
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class OmniLLM:
|
|
"""Multi-process pipelined OmniLLM.
|
|
|
|
- Per-stage process with copy-based IPC (Queues)
|
|
- window=-1 across stages (downstream starts when upstream finishes)
|
|
- max_inflight=1 per stage (serial within a stage), but pipeline across stages
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
model: str,
|
|
stage_configs_path: Optional[str] = None,
|
|
log_stats: bool = False,
|
|
log_file: Optional[str] = None,
|
|
init_sleep_seconds: int = 20,
|
|
shm_threshold_bytes: int = 65536,
|
|
batch_timeout: int = 10,
|
|
init_timeout: int = 300,
|
|
**kwargs,
|
|
):
|
|
self.batch_timeout = batch_timeout
|
|
self._enable_stats: bool = bool(log_stats)
|
|
# Do NOT call super().__init__ to avoid creating OmniStageLLM instances in parent.
|
|
if stage_configs_path is None:
|
|
self.stage_configs = load_stage_configs_from_model(model)
|
|
else:
|
|
self.stage_configs = load_stage_configs_from_yaml(stage_configs_path)
|
|
|
|
# Optional file handler for orchestrator
|
|
self._log_file = log_file
|
|
if self._log_file:
|
|
remove_old_logs(self._log_file, len(self.stage_list))
|
|
configure_orchestrator_logger(logger, self._log_file)
|
|
|
|
self._stats_file, self._overall_stats_file = init_stats_paths(self._enable_stats, self._log_file)
|
|
self._initialize_stages(model, init_sleep_seconds, shm_threshold_bytes, init_timeout)
|
|
|
|
def _initialize_stages(
|
|
self,
|
|
model: str,
|
|
init_sleep_seconds: int,
|
|
shm_threshold_bytes: int,
|
|
init_timeout: int,
|
|
) -> None:
|
|
self.stage_list: list[OmniStage] = []
|
|
|
|
# Build OmniStage instances in parallel, preserve original order
|
|
def _build_stage(idx_cfg: tuple[int, Any]) -> tuple[int, OmniStage]:
|
|
idx, cfg = idx_cfg
|
|
return idx, OmniStage(cfg)
|
|
|
|
with ThreadPoolExecutor(max_workers=min(len(self.stage_configs), max(1, os.cpu_count() or 1))) as executor:
|
|
futures = [executor.submit(_build_stage, (idx, cfg)) for idx, cfg in enumerate(self.stage_configs)]
|
|
results: list[tuple[int, OmniStage]] = []
|
|
for fut in as_completed(futures):
|
|
results.append(fut.result())
|
|
results.sort(key=lambda x: x[0])
|
|
self.stage_list = [st for _, st in results]
|
|
logger.debug("[Orchestrator] Loaded %d stages", len(self.stage_list))
|
|
|
|
self._ctx = mp.get_context("spawn")
|
|
self._stage_in_queues: list[mp.Queue] = []
|
|
self._stage_out_queues: list[mp.Queue] = []
|
|
self._init_sleep_seconds = max(0, int(init_sleep_seconds))
|
|
self._shm_threshold_bytes = max(0, int(shm_threshold_bytes))
|
|
self._start_stage_processes(model)
|
|
# Wait for all stages to report readiness before seeding
|
|
self._stages_ready: set[int] = set()
|
|
self._wait_for_stages_ready(timeout=init_timeout)
|
|
|
|
def _start_stage_processes(self, model: str) -> None:
|
|
for stage_id, stage in enumerate(self.stage_list):
|
|
# Use unbounded queues to avoid deadlock when seeding many requests
|
|
in_q: mp.Queue = self._ctx.Queue(maxsize=0)
|
|
out_q: mp.Queue = self._ctx.Queue(maxsize=0)
|
|
self._stage_in_queues.append(in_q)
|
|
self._stage_out_queues.append(out_q)
|
|
|
|
# Attach queues and start Stage-owned worker process
|
|
stage.attach_queues(in_q, out_q)
|
|
stage.init_stage_worker(
|
|
model,
|
|
log_file=self._log_file,
|
|
shm_threshold_bytes=self._shm_threshold_bytes,
|
|
ctx=self._ctx,
|
|
batch_timeout=self.batch_timeout,
|
|
)
|
|
logger.debug("[Orchestrator] Stage-%s process started", stage_id)
|
|
time.sleep(self._init_sleep_seconds)
|
|
|
|
def close(self) -> None:
|
|
for q in self._stage_in_queues:
|
|
try:
|
|
q.put_nowait(None)
|
|
except Exception as e:
|
|
logger.warning(
|
|
"[Orchestrator] Failed to send shutdown signal to stage input queue: %s",
|
|
e,
|
|
)
|
|
for stage in self.stage_list:
|
|
try:
|
|
stage.stop_stage_worker()
|
|
except Exception as e:
|
|
logger.warning("[Orchestrator] Failed to stop stage worker: %s", e)
|
|
|
|
def __del__(self) -> None: # best-effort
|
|
try:
|
|
self.close()
|
|
except Exception as e:
|
|
logger.debug("[Orchestrator] __del__ close() raised: %s", e, exc_info=True)
|
|
|
|
def generate(
|
|
self,
|
|
prompts: Union[PromptType, Sequence[PromptType]],
|
|
sampling_params_list: Optional[Union[SamplingParams, Sequence[SamplingParams]]] = None,
|
|
) -> list[OmniRequestOutput]:
|
|
try:
|
|
return self._run_generation(prompts, sampling_params_list)
|
|
except Exception as e:
|
|
logger.exception("[Orchestrator] Failed to run generation: %s", e)
|
|
raise e
|
|
finally:
|
|
self.close()
|
|
|
|
def _run_generation(
|
|
self,
|
|
prompts: Union[PromptType, Sequence[PromptType]],
|
|
sampling_params_list: Optional[Union[SamplingParams, Sequence[SamplingParams]]] = None,
|
|
) -> list[OmniRequestOutput]:
|
|
logger.debug("[Orchestrator] generate() called")
|
|
if sampling_params_list is None:
|
|
raise ValueError("sampling_params_list is required for pipelined generation")
|
|
if len(sampling_params_list) != len(self.stage_list):
|
|
raise ValueError(f"Expected {len(self.stage_list)} sampling params, got {len(sampling_params_list)}")
|
|
|
|
# Normalize prompts to a list for per-request iteration
|
|
if not isinstance(prompts, (list, tuple)):
|
|
request_prompts: list[PromptType] = [prompts]
|
|
else:
|
|
request_prompts = list(prompts)
|
|
|
|
final_outputs: list[OmniRequestOutput] = []
|
|
|
|
# Orchestrator keeps stage objects for input derivation
|
|
num_stages = len(self.stage_list)
|
|
|
|
# Map from request_id to original prompt
|
|
request_id_to_prompt: dict[int, PromptType] = {i: p for i, p in enumerate(request_prompts)}
|
|
|
|
# Track per-request start time for end-to-end timing
|
|
_req_start_ts: dict[int, float] = {}
|
|
_wall_start_ts: float = time.time()
|
|
|
|
# Determine the final stage for E2E stats (highest stage_id with final_output=True; fallback to last stage)
|
|
final_stage_id_for_e2e = -1
|
|
try:
|
|
for _sid, _st in enumerate(self.stage_list):
|
|
if getattr(_st, "final_output", False):
|
|
final_stage_id_for_e2e = max(final_stage_id_for_e2e, _sid)
|
|
if final_stage_id_for_e2e < 0:
|
|
final_stage_id_for_e2e = len(self.stage_list) - 1
|
|
except Exception as e:
|
|
logger.debug(
|
|
"[Orchestrator] Failed to determine final stage for E2E; falling back to last: %s",
|
|
e,
|
|
exc_info=True,
|
|
)
|
|
final_stage_id_for_e2e = len(self.stage_list) - 1
|
|
# Metrics/aggregation helper
|
|
metrics = OrchestratorMetrics(
|
|
num_stages,
|
|
self._enable_stats,
|
|
self._stats_file,
|
|
self._overall_stats_file,
|
|
_wall_start_ts,
|
|
)
|
|
|
|
# Seed stage-0 queue with all requests
|
|
logger.debug("[Orchestrator] Seeding %d requests into stage-0", len(request_prompts))
|
|
# Mark first input time for stage-0
|
|
metrics.stage_first_ts[0] = metrics.stage_first_ts[0] or time.time()
|
|
|
|
for req_id, prompt in request_id_to_prompt.items():
|
|
sp0: SamplingParams = sampling_params_list[0] # type: ignore[index]
|
|
task = {
|
|
"request_id": req_id,
|
|
"engine_inputs": prompt,
|
|
"sampling_params": sp0,
|
|
}
|
|
self.stage_list[0].submit(task)
|
|
_req_start_ts[req_id] = time.time()
|
|
logger.debug("[Orchestrator] Enqueued request %s to stage-0", req_id)
|
|
|
|
# For each stage, forward results to next stage; collect finals at the end
|
|
# We pipeline by continually polling output queues in stage order
|
|
remaining_by_stage: list[int] = [len(request_prompts)] + [0] * (num_stages - 1)
|
|
completed_requests = 0
|
|
total_requests = len(request_prompts)
|
|
|
|
logger.debug(
|
|
"[Orchestrator] Entering scheduling loop: total_requests=%d, stages=%d",
|
|
total_requests,
|
|
num_stages,
|
|
)
|
|
while completed_requests < total_requests:
|
|
made_progress = False
|
|
for stage_id, stage in enumerate(self.stage_list):
|
|
result = stage.try_collect()
|
|
if result is None:
|
|
continue
|
|
|
|
made_progress = True
|
|
req_id = result.get("request_id")
|
|
if "error" in result:
|
|
logger.error(
|
|
"Stage %s error on request %s: %s",
|
|
stage_id,
|
|
req_id,
|
|
result["error"],
|
|
)
|
|
continue
|
|
|
|
if result.get("type") == "stage_ready":
|
|
# Only happens when stage is initialized slower than expected,
|
|
# so we wait for a short time and try again
|
|
time.sleep(0.05)
|
|
continue
|
|
|
|
engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm")
|
|
# Mark last output time for this stage whenever we receive outputs
|
|
metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, time.time())
|
|
try:
|
|
_m = result.get("metrics")
|
|
if _m is not None:
|
|
metrics.on_stage_metrics(stage_id, req_id, _m)
|
|
except Exception as e:
|
|
logger.exception(
|
|
"[Orchestrator] Failed to process metrics for stage %s, req %s: %s",
|
|
stage_id,
|
|
req_id,
|
|
e,
|
|
)
|
|
logger.debug(
|
|
"[Orchestrator] Stage-%s completed request %s; forwarding or finalizing",
|
|
stage_id,
|
|
req_id,
|
|
)
|
|
stage.set_engine_outputs(engine_outputs)
|
|
|
|
if getattr(stage, "final_output", False):
|
|
final_outputs.append(
|
|
OmniRequestOutput(
|
|
stage_id=stage_id,
|
|
final_output_type=stage.final_output_type, # type: ignore[attr-defined]
|
|
request_output=engine_outputs,
|
|
)
|
|
)
|
|
logger.debug(
|
|
"[Orchestrator] Request %s finalized at stage-%s",
|
|
req_id,
|
|
stage_id,
|
|
)
|
|
|
|
# End-to-end timing and time-per-token for final output
|
|
# (only once per request at the designated final stage)
|
|
try:
|
|
rid_int = int(req_id) if isinstance(req_id, (int, str)) and str(req_id).isdigit() else req_id
|
|
if stage_id == final_stage_id_for_e2e and rid_int not in metrics.e2e_done:
|
|
metrics.on_finalize_request(
|
|
stage_id,
|
|
req_id,
|
|
engine_outputs,
|
|
_req_start_ts.get(req_id, _wall_start_ts),
|
|
)
|
|
except Exception as e:
|
|
logger.exception(
|
|
"[Orchestrator] Finalize request handling error for req %s at stage %s: %s",
|
|
req_id,
|
|
stage_id,
|
|
e,
|
|
)
|
|
|
|
next_stage_id = stage_id + 1
|
|
if next_stage_id < num_stages:
|
|
next_stage: OmniStage = self.stage_list[next_stage_id]
|
|
next_inputs = next_stage.process_engine_inputs(self.stage_list, [request_id_to_prompt[req_id]])
|
|
sp_next: SamplingParams = sampling_params_list[next_stage_id] # type: ignore[index]
|
|
try:
|
|
# Measure transfer size and time (encode + enqueue)
|
|
size_bytes = 0
|
|
try:
|
|
size_bytes = len(_set(next_inputs))
|
|
except Exception:
|
|
size_bytes = 0
|
|
t0 = time.time()
|
|
ipc_payload = _encode(
|
|
next_inputs,
|
|
getattr(self, "_shm_threshold_bytes", 65536),
|
|
obj_key="engine_inputs",
|
|
shm_key="engine_inputs_shm",
|
|
)
|
|
ipc_payload.update(
|
|
{
|
|
"request_id": req_id,
|
|
"sampling_params": sp_next,
|
|
"sent_ts": time.time(),
|
|
}
|
|
)
|
|
self.stage_list[next_stage_id].submit(ipc_payload)
|
|
t1 = time.time()
|
|
tx_ms = (t1 - t0) * 1000.0
|
|
metrics.on_forward(
|
|
stage_id,
|
|
next_stage_id,
|
|
req_id,
|
|
int(size_bytes),
|
|
float(tx_ms),
|
|
bool("engine_inputs_shm" in ipc_payload),
|
|
)
|
|
except Exception as e:
|
|
logger.warning(
|
|
"[Orchestrator] IPC encode failed for req %s: %s; falling back to inline payload",
|
|
req_id,
|
|
e,
|
|
)
|
|
self.stage_list[next_stage_id].submit(
|
|
{
|
|
"request_id": req_id,
|
|
"engine_inputs": next_inputs,
|
|
"sampling_params": sp_next,
|
|
}
|
|
)
|
|
logger.debug(
|
|
"[Orchestrator] Forwarded request %s to stage-%s",
|
|
req_id,
|
|
next_stage_id,
|
|
)
|
|
remaining_by_stage[next_stage_id] += 1
|
|
else:
|
|
completed_requests += 1
|
|
logger.debug(
|
|
"[Orchestrator] Request %s fully completed (%d/%d)",
|
|
req_id,
|
|
completed_requests,
|
|
total_requests,
|
|
)
|
|
|
|
if not made_progress:
|
|
time.sleep(0.005)
|
|
logger.debug("[Orchestrator] All requests completed")
|
|
|
|
# Summarize and print stats
|
|
try:
|
|
summary = metrics.build_and_log_summary(final_stage_id_for_e2e)
|
|
logger.info("[Summary] %s", summary)
|
|
except Exception as e:
|
|
logger.exception("[Orchestrator] Failed to build/log summary: %s", e)
|
|
|
|
return final_outputs
|
|
|
|
def _wait_for_stages_ready(self, timeout: int = 120) -> None:
|
|
deadline = time.time() + max(0, int(timeout))
|
|
num_stages = len(self.stage_list)
|
|
while len(self._stages_ready) < num_stages and time.time() < deadline:
|
|
progressed = False
|
|
for stage_id, stage in enumerate(self.stage_list):
|
|
if stage_id in self._stages_ready:
|
|
continue
|
|
result = stage.try_collect()
|
|
if result is None:
|
|
continue
|
|
progressed = True
|
|
if result.get("type") == "stage_ready":
|
|
self._stages_ready.add(stage_id)
|
|
logger.debug("[Orchestrator] Stage-%s reported ready", stage_id)
|
|
else:
|
|
# No user data should arrive before seeding; ignore other messages
|
|
pass
|
|
if not progressed:
|
|
time.sleep(0.01)
|
|
if len(self._stages_ready) < num_stages:
|
|
not_ready = sorted(set(range(num_stages)) - set(self._stages_ready))
|
|
logger.warning(
|
|
"[Orchestrator] Initialization timeout: only %s/%s stages are ready; not ready: %s",
|
|
len(self._stages_ready),
|
|
num_stages,
|
|
not_ready,
|
|
)
|
|
# Provide actionable suggestions before shutdown
|
|
try:
|
|
suggestions = [
|
|
"Verify GPU/device assignment in config (runtime.devices) is correct.",
|
|
"Check GPU/host memory availability; reduce model or batch size if needed.",
|
|
"Check model weights path and network reachability (if loading remotely).",
|
|
"Increase initialization wait time (init_sleep_seconds or call-site timeout).",
|
|
]
|
|
if getattr(self, "_log_file", None):
|
|
suggestions.append(f"Inspect per-stage log files for details: {self._log_file}.stage<id>.log")
|
|
logger.error(
|
|
"[Orchestrator] Stage initialization failed, shutting down. Suggestions:\n- %s",
|
|
"\n- ".join(suggestions),
|
|
)
|
|
except Exception:
|
|
# Best-effort logging of suggestions
|
|
logger.error(
|
|
"[Orchestrator] Stage initialization failed and an error occurred while logging suggestions",
|
|
)
|
|
|
|
# Attempt graceful shutdown of all stages before exiting
|
|
try:
|
|
self.close()
|
|
except Exception:
|
|
pass
|
|
|
|
# Terminate the current process with non-zero exit code
|
|
try:
|
|
sys.exit(1)
|
|
except SystemExit:
|
|
raise
|
|
except Exception:
|
|
os._exit(1)
|
|
|
|
|
|
class OmniStageLLM(LLM):
|
|
def __init__(
|
|
self,
|
|
model: str,
|
|
compilation_config: Optional[Union[int, dict[str, Any], CompilationConfig]] = None,
|
|
hf_overrides: Optional[HfOverrides] = None,
|
|
structured_outputs_config: Optional[Union[dict[str, Any], StructuredOutputsConfig]] = None,
|
|
**kwargs,
|
|
):
|
|
"""LLM constructor."""
|
|
if "disable_log_stats" not in kwargs:
|
|
kwargs["disable_log_stats"] = True
|
|
|
|
if "worker_cls" in kwargs:
|
|
worker_cls = kwargs["worker_cls"]
|
|
# if the worker_cls is not qualified string name,
|
|
# we serialize it using cloudpickle to avoid pickling issues
|
|
if isinstance(worker_cls, type):
|
|
kwargs["worker_cls"] = cloudpickle.dumps(worker_cls)
|
|
|
|
if "kv_transfer_config" in kwargs and isinstance(kwargs["kv_transfer_config"], dict):
|
|
from vllm.config.kv_transfer import KVTransferConfig
|
|
|
|
raw_config_dict = kwargs["kv_transfer_config"]
|
|
try:
|
|
kwargs["kv_transfer_config"] = KVTransferConfig(**raw_config_dict)
|
|
except ValidationError as e:
|
|
logger.error(
|
|
"Failed to convert 'kv_transfer_config' dict to KVTransferConfig object. Dict: %s. Error: %s",
|
|
raw_config_dict,
|
|
e,
|
|
)
|
|
# Consider re-raising a more specific vLLM error or ValueError
|
|
# to provide better context to the user.
|
|
raise ValueError(f"Invalid 'kv_transfer_config' provided: {e}") from e
|
|
|
|
if hf_overrides is None:
|
|
hf_overrides = {}
|
|
|
|
if compilation_config is not None:
|
|
if isinstance(compilation_config, int):
|
|
compilation_config_instance = CompilationConfig(level=compilation_config)
|
|
elif isinstance(compilation_config, dict):
|
|
compilation_config_instance = CompilationConfig(
|
|
**{k: v for k, v in compilation_config.items() if is_init_field(CompilationConfig, k)}
|
|
)
|
|
else:
|
|
compilation_config_instance = compilation_config
|
|
else:
|
|
compilation_config_instance = CompilationConfig()
|
|
|
|
if structured_outputs_config is not None:
|
|
if isinstance(structured_outputs_config, dict):
|
|
structured_outputs_instance = StructuredOutputsConfig(
|
|
**{k: v for k, v in structured_outputs_config.items() if is_init_field(StructuredOutputsConfig, k)}
|
|
)
|
|
else:
|
|
structured_outputs_instance = structured_outputs_config
|
|
else:
|
|
structured_outputs_instance = StructuredOutputsConfig()
|
|
|
|
engine_args = OmniEngineArgs(
|
|
model=model,
|
|
hf_overrides=hf_overrides,
|
|
compilation_config=compilation_config_instance,
|
|
structured_outputs_config=structured_outputs_instance,
|
|
**kwargs,
|
|
)
|
|
|
|
# Create the Engine (autoselects V0 vs V1)
|
|
self.llm_engine = LLMEngine.from_engine_args(engine_args=engine_args, usage_context=UsageContext.LLM_CLASS)
|
|
self.llm_engine.output_processor = MultimodalOutputProcessor(
|
|
tokenizer=self.llm_engine.tokenizer,
|
|
log_stats=self.llm_engine.log_stats,
|
|
engine_core_output_type=engine_args.engine_output_type,
|
|
)
|
|
self.llm_engine.processor = OmniProcessor(
|
|
vllm_config=self.llm_engine.vllm_config, tokenizer=self.llm_engine.tokenizer
|
|
)
|
|
self.engine_class = type(self.llm_engine)
|
|
|
|
self.request_counter = Counter()
|
|
self.default_sampling_params: Union[dict[str, Any], None] = None
|
|
|
|
supported_tasks = self.llm_engine.get_supported_tasks() # type: ignore
|
|
|
|
logger.info("Supported_tasks: %s", supported_tasks)
|
|
|
|
self.supported_tasks = supported_tasks
|
|
|
|
# Load the Input/Output processor plugin if any
|
|
io_processor_plugin = self.llm_engine.model_config.io_processor_plugin
|
|
self.io_processor = get_io_processor(self.llm_engine.vllm_config, io_processor_plugin)
|