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

[Feature] logger refactor (#466)

Browse Source 
Signed-off-by: dengyunyang <584797741@qq.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
dengyunyangandHongsheng Liu authored and GitHub committed 2025-12-26 08:42:06 +08:00
1 parent b3d212e985
commit a500ca1f28
5 files changed
+310 -645
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/entrypoints/test_omni_llm.py
+1 -1
View File
@@ -314,7 +314,7 @@ def _setup_log_mocks(monkeypatch):
# Mock OrchestratorMetrics to be a simple class that doesn't require file operations
class _FakeOrchestratorMetrics:
def __init__(self, num_stages, enable_stats, stats_file, overall_stats_file, wall_start_ts):
def __init__(self, num_stages, enable_stats, wall_start_ts):
self.num_stages = num_stages
self.enable_stats = enable_stats
self.stage_first_ts = [None] * num_stages
vllm_omni/entrypoints/async_omni.py
+4 -20
View File
@@ -6,6 +6,8 @@ import time
from argparse import Namespace
from collections.abc import AsyncGenerator, Iterable, Mapping, Sequence
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import asdict
from pprint import pformat
from typing import Any
import torch
@@ -50,9 +52,6 @@ from vllm_omni.engine.output_processor import MultimodalOutputProcessor
from vllm_omni.entrypoints.client_request_state import ClientRequestState
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 maybe_load_from_ipc as _load
@@ -82,7 +81,6 @@ class AsyncOmni(EngineClient):
- stage_configs_path: Optional path to YAML file containing stage
configurations. If None, configurations are loaded from the model.
- log_stats: Whether to enable statistics logging
- log_file: Optional path prefix for log files. If provided, logs will
be written to files with stage-specific suffixes.
- init_sleep_seconds: Number of seconds to sleep between starting
each stage process during initialization
@@ -139,13 +137,7 @@ class AsyncOmni(EngineClient):
self.stage_list: list[OmniStage] = []
self.default_sampling_params_list: list[SamplingParams] = []
# Optional file handler for orchestrator
self._log_file = cli_args.log_file
if self._log_file:
remove_old_logs(self._log_file, len(self.stage_configs))
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, cli_args.init_sleep_seconds, cli_args.shm_threshold_bytes, cli_args.init_timeout)
# Here is a duplicated init for the first stage to get the tokenizer,
@@ -231,7 +223,6 @@ class AsyncOmni(EngineClient):
stage.init_stage_worker(
model,
is_async=True,
log_file=self._log_file,
shm_threshold_bytes=self._shm_threshold_bytes,
ctx=self._ctx if self.worker_backend != "ray" else None,
batch_timeout=self.batch_timeout,
@@ -358,8 +349,6 @@ class AsyncOmni(EngineClient):
metrics = OrchestratorMetrics(
num_stages,
self._enable_stats,
self._stats_file,
self._overall_stats_file,
_wall_start_ts,
)
# Seed stage-0 queue with all requests
@@ -399,7 +388,7 @@ class AsyncOmni(EngineClient):
# 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")
_m = asdict(result.get("metrics"))
if _m is not None:
metrics.on_stage_metrics(stage_id, req_id, _m)
except Exception as e:
@@ -502,7 +491,7 @@ class AsyncOmni(EngineClient):
# Summarize and print stats
try:
summary = metrics.build_and_log_summary(final_stage_id_for_e2e)
logger.info("[Summary] %s", summary)
logger.info("[Summary] %s", pformat(summary, sort_dicts=False))
except Exception as e:
logger.exception("[Orchestrator] Failed to build/log summary: %s", e)
finally:
@@ -559,11 +548,6 @@ class AsyncOmni(EngineClient):
"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",
vllm_omni/entrypoints/log_utils.py
+121 -235
View File
@@ -1,92 +1,16 @@
from __future__ import annotations
import logging
import os
import time
from dataclasses import dataclass
from pprint import pformat
from typing import Any
from vllm_omni.entrypoints.stage_utils import append_jsonl as _append_jsonl
from vllm.logger import init_logger
def remove_old_logs(log_file: str | None, num_stages: int) -> None:
try:
if not log_file:
return
try:
if os.path.exists(log_file):
os.remove(log_file)
except Exception:
pass
# Per-stage logs and stats
for sid in range(num_stages):
try:
p = f"{log_file}.stage{sid}.log"
if os.path.exists(p):
os.remove(p)
except Exception:
pass
try:
p = f"{log_file}.stage{sid}.stats.jsonl"
if os.path.exists(p):
os.remove(p)
except Exception:
pass
# Orchestrator stats files
try:
p = f"{log_file}.orchestrator.stats.jsonl"
if os.path.exists(p):
os.remove(p)
except Exception:
pass
try:
p = f"{log_file}.overall.stats.jsonl"
if os.path.exists(p):
os.remove(p)
except Exception:
pass
except Exception:
pass
def configure_orchestrator_logger(logger: logging.Logger, log_file: str | None) -> None:
try:
if not log_file:
return
has_file_handler = any(isinstance(h, logging.FileHandler) for h in logger.handlers)
if not has_file_handler:
fh = logging.FileHandler(log_file)
fh.setLevel(logging.DEBUG)
fh.setFormatter(logging.Formatter("%(asctime)s [PID:%(process)d] %(levelname)s: %(message)s"))
logger.addHandler(fh)
logger.setLevel(logging.DEBUG)
except Exception:
pass
def init_stats_paths(enable_stats: bool, log_file: str | None) -> tuple[str | None, str | None]:
stats_file: str | None = None
overall_file: str | None = None
try:
if enable_stats and log_file:
stats_file = f"{log_file}.orchestrator.stats.jsonl"
overall_file = f"{log_file}.overall.stats.jsonl"
except Exception:
stats_file = None
overall_file = None
return stats_file, overall_file
def _safe_append_jsonl(path: str | None, record: dict[str, Any]) -> None:
if not path:
return
try:
_append_jsonl(path, record) # type: ignore[arg-type]
except Exception:
pass
logger = init_logger(__name__)
def log_transfer_tx(
stats_file: str | None,
from_stage: int,
to_stage: int,
request_id: Any,
@@ -94,23 +18,24 @@ def log_transfer_tx(
tx_time_ms: float,
used_shm: bool,
) -> None:
_safe_append_jsonl(
stats_file,
{
"type": "transfer_stats",
"from_stage": from_stage,
"to_stage": to_stage,
"request_id": request_id,
"size_bytes": int(size_bytes),
"tx_time_ms": float(tx_time_ms),
"tx_mbps": (float(size_bytes) * 8.0) / (max(tx_time_ms, 1e-6) * 1000.0),
"used_shm": bool(used_shm),
},
logger.info(
pformat(
{
"type": "transfer_stats",
"from_stage": from_stage,
"to_stage": to_stage,
"request_id": request_id,
"size_bytes": int(size_bytes),
"tx_time_ms": float(tx_time_ms),
"tx_mbps": (float(size_bytes) * 8.0) / (max(tx_time_ms, 1e-6) * 1000.0),
"used_shm": bool(used_shm),
},
sort_dicts=False,
)
)
def log_transfer_rx(
stats_file: str | None,
from_stage: int,
to_stage: int,
request_id: Any,
@@ -118,25 +43,26 @@ def log_transfer_rx(
rx_decode_time_ms: float,
in_flight_time_ms: float,
) -> None:
_safe_append_jsonl(
stats_file,
{
"type": "transfer_rx_stats",
"from_stage": from_stage,
"to_stage": to_stage,
"request_id": request_id,
"rx_bytes": int(rx_bytes),
"rx_decode_time_ms": float(rx_decode_time_ms),
"in_flight_time_ms": float(in_flight_time_ms),
"rx_time_per_kb_ms": (
(float(rx_decode_time_ms) / max(float(rx_bytes) / 1024.0, 1e-6)) if rx_bytes > 0 else 0.0
),
},
logger.info(
pformat(
{
"type": "transfer_rx_stats",
"from_stage": from_stage,
"to_stage": to_stage,
"request_id": request_id,
"rx_bytes": int(rx_bytes),
"rx_decode_time_ms": float(rx_decode_time_ms),
"in_flight_time_ms": float(in_flight_time_ms),
"rx_time_per_kb_ms": (
(float(rx_decode_time_ms) / max(float(rx_bytes) / 1024.0, 1e-6)) if rx_bytes > 0 else 0.0
),
},
sort_dicts=False,
)
)
def log_transfer_total(
stats_file: str | None,
from_stage: int,
to_stage: int,
request_id: Any,
@@ -146,60 +72,28 @@ def log_transfer_total(
rx_decode_time_ms: float,
total_time_ms: float,
) -> None:
_safe_append_jsonl(
stats_file,
{
"type": "transfer_total_stats",
"from_stage": from_stage,
"to_stage": to_stage,
"request_id": request_id,
"size_bytes": int(size_bytes),
"tx_time_ms": float(tx_time_ms),
"in_flight_time_ms": float(in_flight_time_ms),
"rx_decode_time_ms": float(rx_decode_time_ms),
"total_time_ms": float(total_time_ms),
"total_time_per_kb_ms": (
float(total_time_ms) / max(float(size_bytes) / 1024.0, 1e-6) if size_bytes > 0 else 0.0
),
},
logger.info(
pformat(
{
"type": "transfer_total_stats",
"from_stage": from_stage,
"to_stage": to_stage,
"request_id": request_id,
"size_bytes": int(size_bytes),
"tx_time_ms": float(tx_time_ms),
"in_flight_time_ms": float(in_flight_time_ms),
"rx_decode_time_ms": float(rx_decode_time_ms),
"total_time_ms": float(total_time_ms),
"total_time_per_kb_ms": (
float(total_time_ms) / max(float(size_bytes) / 1024.0, 1e-6) if size_bytes > 0 else 0.0
),
},
sort_dicts=False,
)
)
def log_orchestrator_e2e(
stats_file: str | None,
request_id: Any,
final_stage_id: int,
e2e_time_ms: float,
num_tokens_out: int,
e2e_time_per_token_ms: float,
) -> None:
_safe_append_jsonl(
stats_file,
{
"type": "orchestrator_request_e2e",
"request_id": request_id,
"final_stage_id": final_stage_id,
"e2e_time_ms": float(e2e_time_ms),
"num_tokens_out": int(num_tokens_out),
"e2e_time_per_token_ms": float(e2e_time_per_token_ms),
},
)
def log_orchestrator_summary(stats_file: str | None, summary: dict[str, Any]) -> None:
_safe_append_jsonl(stats_file, {"type": "orchestrator_summary", **summary})
def log_overall_summary(overall_stats_file: str | None, summary: dict[str, Any]) -> None:
_safe_append_jsonl(overall_stats_file, {"type": "overall_summary", **summary})
def log_overall_record(overall_stats_file: str | None, record: dict[str, Any]) -> None:
_safe_append_jsonl(overall_stats_file, record)
def log_stage_request_stats(
stats_file: str | None,
stage_id: int,
request_id: Any,
batch_size: int,
@@ -210,85 +104,46 @@ def log_stage_request_stats(
rx_decode_time_ms: float,
rx_mbps: float,
) -> None:
_safe_append_jsonl(
stats_file,
{
"type": "stage_request_stats",
"stage_id": stage_id,
"request_id": request_id,
"batch_size": int(batch_size),
"num_tokens_out": int(num_tokens_out),
"stage_gen_time_ms": float(stage_gen_time_ms),
"tokens_per_s": float(tokens_per_s),
"rx_transfer_bytes": int(rx_transfer_bytes),
"rx_decode_time_ms": float(rx_decode_time_ms),
"rx_mbps": float(rx_mbps),
},
)
def log_stage_running_avg(
stats_file: str | None,
stage_id: int,
total_tokens: int,
total_gen_time_ms: float,
avg_tokens_per_s: float,
) -> None:
_safe_append_jsonl(
stats_file,
{
"type": "stage_running_avg",
"stage_id": stage_id,
"total_tokens": int(total_tokens),
"total_gen_time_ms": float(total_gen_time_ms),
"avg_tokens_per_s": float(avg_tokens_per_s),
},
)
def log_stage_batch_stats(
stats_file: str | None,
stage_id: int,
batch_size: int,
batch_gen_time_ms: float,
request_ids: list[Any],
) -> None:
_safe_append_jsonl(
stats_file,
{
"type": "stage_batch_stats",
"stage_id": stage_id,
"batch_size": int(batch_size),
"batch_gen_time_ms": float(batch_gen_time_ms),
"request_ids": list(request_ids),
},
logger.info(
pformat(
{
"type": "Request_stage_stats",
"stage_id": stage_id,
"request_id": request_id,
"batch_size": int(batch_size),
"num_tokens_out": int(num_tokens_out),
"stage_gen_time_ms": float(stage_gen_time_ms),
"tokens_per_s": float(tokens_per_s),
"rx_transfer_bytes": int(rx_transfer_bytes),
"rx_decode_time_ms": float(rx_decode_time_ms),
"rx_mbps": float(rx_mbps),
},
sort_dicts=False,
)
)
def compute_and_log_stage_request_stats(
stats_file: str | None,
stage_id: int,
request_id: Any,
batch_size: int,
engine_outputs: list[Any],
num_engine_outputs: int,
stage_gen_time_ms: float,
rx_transfer_bytes: int,
rx_decode_time_ms: float,
) -> None:
"""Compute per-request metrics and log them in one call."""
num_tokens = count_tokens_from_outputs(engine_outputs)
tokens_per_s = (num_tokens * 1000.0 / stage_gen_time_ms) if stage_gen_time_ms > 0 else 0.0
tokens_per_s = (num_engine_outputs * 1000.0 / stage_gen_time_ms) if stage_gen_time_ms > 0 else 0.0
rx_mbps = (
(float(rx_transfer_bytes) * 8.0) / (max(float(rx_decode_time_ms), 1e-6) * 1000.0)
if rx_transfer_bytes > 0
else 0.0
)
log_stage_request_stats(
stats_file,
stage_id,
request_id,
int(batch_size),
int(num_tokens),
int(num_engine_outputs),
float(stage_gen_time_ms),
float(tokens_per_s),
int(rx_transfer_bytes),
@@ -489,19 +344,34 @@ def build_transfer_summary(
return summary
@dataclass
class StageStats:
total_token: int
total_gen_time: float
@dataclass
class StageRequestMetrics:
num_tokens_out: int
stage_gen_time_ms: float
batch_id: int
batch_size: int
rx_decode_time_ms: float
rx_transfer_bytes: int
rx_in_flight_time_ms: float
stage_stats: StageStats
class OrchestratorMetrics:
def __init__(
self,
num_stages: int,
enable_stats: bool,
stats_file: str | None,
overall_stats_file: str | None,
wall_start_ts: float,
) -> None:
self.num_stages = int(num_stages)
self.enable_stats = bool(enable_stats)
self.stats_file = stats_file
self.overall_stats_file = overall_stats_file
self.stage_total_time_ms: list[float] = [0.0 for _ in range(self.num_stages)]
self.stage_total_tokens: list[int] = [0 for _ in range(self.num_stages)]
self.stage_req_counts: list[int] = [0 for _ in range(self.num_stages)]
@@ -529,6 +399,32 @@ class OrchestratorMetrics:
req_id,
metrics,
)
if self.enable_stats:
compute_and_log_stage_request_stats(
stage_id=stage_id,
request_id=req_id,
batch_size=metrics.get("batch_size"),
num_engine_outputs=metrics.get("num_tokens_out"),
stage_gen_time_ms=metrics.get("stage_gen_time_ms"),
rx_decode_time_ms=metrics.get("rx_decode_time_ms"),
rx_transfer_bytes=metrics.get("rx_transfer_bytes"),
)
if stage_stats := metrics.get("stage_stats", None):
total_token = int(stage_stats.get("total_token"))
total_gen_time = float(stage_stats.get("total_gen_time"))
_avg_tokens_per_s = (total_token * 1000.0 / total_gen_time) if total_gen_time > 0 else 0.0
logger.info(
pformat(
{
"type": "Stage_running_avg",
"stage_id": stage_id,
"total_tokens": total_token,
"total_gen_time_ms": total_gen_time,
"avg_tokens_per_s": _avg_tokens_per_s,
},
sort_dicts=False,
)
)
try:
batch_id_raw = metrics.get("batch_id", None)
if batch_id_raw is not None:
@@ -551,9 +447,8 @@ class OrchestratorMetrics:
rx_ms,
in_flight_ms,
)
if self.enable_stats and self.stats_file and stage_id > 0:
if self.enable_stats and stage_id > 0:
log_transfer_rx(
self.stats_file,
stage_id - 1,
stage_id,
req_id,
@@ -564,7 +459,6 @@ class OrchestratorMetrics:
if combined is not None:
size_b_c, tx_ms_c, total_ms_c = combined
log_transfer_total(
self.stats_file,
stage_id - 1,
stage_id,
req_id,
@@ -587,9 +481,8 @@ class OrchestratorMetrics:
# Mark first input time for the destination stage if not set
if self.stage_first_ts[to_stage] is None:
self.stage_first_ts[to_stage] = time.time()
if self.enable_stats and self.stats_file:
if self.enable_stats:
log_transfer_tx(
self.stats_file,
from_stage,
to_stage,
req_id,
@@ -623,20 +516,17 @@ class OrchestratorMetrics:
self.e2e_done.add(rid_key)
pr = self.per_request.setdefault(rid_key, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
per_req_record = {
"type": "overall_request",
"type": "request_level_metrics",
"request_id": rid_key,
"e2e_time_ms": e2e_ms,
"e2e_tpt": (e2e_ms / num_tokens) if num_tokens > 0 else 0.0,
"num_tokens_out": int(num_tokens),
"transfers_total_time_ms": float(pr.get("transfers_ms", 0.0)),
"transfers_total_bytes": int(pr.get("transfers_bytes", 0)),
"stages": pr.get("stages", {}),
}
self.sum_per_request_transfer_ms += float(pr.get("transfers_ms", 0.0))
if self.enable_stats and self.overall_stats_file:
log_overall_record(self.overall_stats_file, per_req_record)
if self.enable_stats and self.stats_file:
e2e_tpt = (e2e_ms / num_tokens) if num_tokens > 0 else 0.0
log_orchestrator_e2e(self.stats_file, req_id, stage_id, e2e_ms, int(num_tokens), e2e_tpt)
logger.info(pformat(per_req_record, sort_dicts=False))
def build_and_log_summary(self, final_stage_id_to_prompt: dict[str, int]) -> dict[str, Any]:
# Compute stage summary using wall time between first input and last output per stage
@@ -677,8 +567,4 @@ class OrchestratorMetrics:
"stages": stage_summary,
"transfers": transfer_summary,
}
if self.enable_stats and self.stats_file:
log_orchestrator_summary(self.stats_file, summary)
if self.enable_stats and self.overall_stats_file:
log_overall_summary(self.overall_stats_file, summary)
return summary
vllm_omni/entrypoints/omni_llm.py
+4 -20
View File
@@ -4,6 +4,8 @@ import time
import uuid
from collections.abc import Sequence
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import asdict
from pprint import pformat
from typing import Any
import cloudpickle
@@ -37,9 +39,6 @@ from vllm_omni.engine.input_processor import OmniInputProcessor
from vllm_omni.engine.output_processor import MultimodalOutputProcessor
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 maybe_load_from_ipc as _load
@@ -69,8 +68,6 @@ class OmniLLM:
stage_configs_path: Optional path to YAML file containing stage
configurations. If None, configurations are loaded from the model.
log_stats: Whether to enable statistics logging
log_file: Optional path prefix for log files. If provided, logs will
be written to files with stage-specific suffixes.
init_sleep_seconds: Number of seconds to sleep between starting
each stage process during initialization
shm_threshold_bytes: Threshold in bytes for using shared memory
@@ -92,7 +89,6 @@ class OmniLLM:
model: str,
stage_configs_path: str | None = None,
log_stats: bool = False,
log_file: str | None = None,
init_sleep_seconds: int = 20,
shm_threshold_bytes: int = 65536,
batch_timeout: int = 10,
@@ -118,14 +114,7 @@ class OmniLLM:
self.config_path, worker_backend=self.worker_backend, shm_threshold_bytes=shm_threshold_bytes
)
# Optional file handler for orchestrator
self._log_file = 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)
if self._log_file:
remove_old_logs(self._log_file, len(self.stage_list))
configure_orchestrator_logger(logger, self._log_file)
def _initialize_stages(
self,
@@ -187,7 +176,6 @@ class OmniLLM:
stage.init_stage_worker(
model,
log_file=self._log_file,
shm_threshold_bytes=self._shm_threshold_bytes,
ctx=self._ctx if self.worker_backend != "ray" else None,
batch_timeout=self.batch_timeout,
@@ -308,8 +296,6 @@ class OmniLLM:
metrics = OrchestratorMetrics(
num_stages,
self._enable_stats,
self._stats_file,
self._overall_stats_file,
_wall_start_ts,
)
@@ -371,7 +357,7 @@ class OmniLLM:
# 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")
_m = asdict(result.get("metrics"))
if _m is not None:
metrics.on_stage_metrics(stage_id, req_id, _m)
except Exception as e:
@@ -480,7 +466,7 @@ class OmniLLM:
# Summarize and print stats
try:
summary = metrics.build_and_log_summary(final_stage_id_to_prompt)
logger.info("[Summary] %s", summary)
logger.info("[Summary] %s", pformat(summary, sort_dicts=False))
except Exception as e:
logger.exception("[Orchestrator] Failed to build/log summary: %s", e)
@@ -522,8 +508,6 @@ class OmniLLM:
"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),
vllm_omni/entrypoints/omni_stage.py
+180 -369
View File
@@ -13,7 +13,6 @@ the original input processing utilities for cross-stage data wiring.
import asyncio
import fcntl
import importlib
import logging
import multiprocessing as mp
import os
import queue
@@ -87,9 +86,7 @@ class OmniStage:
self._in_q: mp.Queue | None = None
self._out_q: mp.Queue | None = None
self._proc: mp.Process | None = None
self._log_file: str | None = None
self._shm_threshold_bytes: int = 65536
self._logger = logging.getLogger(__name__)
def set_engine(self, engine: LLMEngine) -> None:
"""Set the LLM engine for this stage.
@@ -163,7 +160,6 @@ class OmniStage:
model: str,
*,
is_async: bool = False,
log_file: str | None = None,
shm_threshold_bytes: int = 65536,
ctx: mp.context.BaseContext | None = None,
batch_timeout: int = 10,
@@ -179,7 +175,6 @@ class OmniStage:
Args:
model: Model name or path to load
is_async: Whether to use async engine (default: False)
log_file: Optional log file path prefix for stage-specific logs
shm_threshold_bytes: Threshold for using shared memory for IPC
ctx: Optional multiprocessing context (default: spawn)
batch_timeout: Timeout in seconds for batching requests
@@ -191,7 +186,6 @@ class OmniStage:
AssertionError: If queues are not attached before calling this method
"""
assert self._in_q is not None and self._out_q is not None, "Queues must be attached before start_process"
self._log_file = log_file
if worker_backend == "ray":
ray_placement_group = kwargs.get("ray_placement_group", None)
@@ -211,54 +205,60 @@ class OmniStage:
"shm_threshold_bytes": self._shm_threshold_bytes,
"connectors_config": connectors_config or {},
}
if worker_backend == "ray":
if is_async:
self._ray_actor = start_ray_actor(
_stage_worker_async_entry,
ray_placement_group,
self.stage_id,
self,
model=model,
stage_payload=stage_payload,
batch_timeout=batch_timeout,
)
else:
self._ray_actor = start_ray_actor(
_stage_worker,
ray_placement_group,
self.stage_id,
model=model,
stage_payload=stage_payload,
in_q=self._in_q,
out_q=self._out_q,
log_file=self._log_file,
batch_timeout=batch_timeout,
)
else:
if is_async:
self._proc = ctx.Process(
target=_stage_worker_async_entry,
args=(
try:
old_env = os.environ.get("VLLM_LOGGING_PREFIX")
new_env = f"[Stage-{self.stage_id}] {'' if old_env is None else old_env}"
os.environ["VLLM_LOGGING_PREFIX"] = new_env
if worker_backend == "ray":
if is_async:
self._ray_actor = start_ray_actor(
_stage_worker_async_entry,
ray_placement_group,
self.stage_id,
self,
model,
stage_payload,
batch_timeout,
),
)
model=model,
stage_payload=stage_payload,
batch_timeout=batch_timeout,
)
else:
self._ray_actor = start_ray_actor(
_stage_worker,
ray_placement_group,
self.stage_id,
model=model,
stage_payload=stage_payload,
in_q=self._in_q,
out_q=self._out_q,
batch_timeout=batch_timeout,
)
else:
self._proc = ctx.Process(
target=_stage_worker,
args=(
model,
stage_payload,
self._in_q,
self._out_q,
self._log_file,
batch_timeout,
),
)
self._proc.start()
if is_async:
self._proc = ctx.Process(
target=_stage_worker_async_entry,
args=(
self,
model,
stage_payload,
batch_timeout,
),
)
else:
self._proc = ctx.Process(
target=_stage_worker,
args=(
model,
stage_payload,
self._in_q,
self._out_q,
batch_timeout,
),
)
self._proc.start()
finally:
if old_env is None:
os.environ.pop("VLLM_LOGGING_PREFIX", None)
else:
os.environ["VLLM_LOGGING_PREFIX"] = old_env
def stop_stage_worker(self) -> None:
"""Stop the stage worker process gracefully.
@@ -271,7 +271,7 @@ class OmniStage:
try:
self._in_q.put_nowait(None)
except Exception as e:
self._logger.warning("[Stage-%s] Failed to send shutdown to in_q: %s", self.stage_id, e)
logger.warning("Failed to send shutdown to in_q: %s", e)
if hasattr(self, "_ray_actor") and self._ray_actor:
kill_ray_actor(self._ray_actor)
@@ -280,21 +280,12 @@ class OmniStage:
try:
self._proc.join(timeout=5)
except Exception as e:
self._logger.debug("[Stage-%s] join() failed: %s", self.stage_id, e, exc_info=True)
logger.debug("join() failed: %s", e)
if self._proc.is_alive():
try:
self._proc.terminate()
except Exception as e:
self._logger.warning("[Stage-%s] terminate() failed: %s", self.stage_id, e)
# Cleanup temporary stage log if we created one (only when no log_file provided)
try:
if not self._log_file:
tmp_path = f"/tmp/omni_stage{self.stage_id}.log"
if os.path.exists(tmp_path):
os.remove(tmp_path)
except Exception:
pass
logger.warning("terminate() failed: %s", e)
def submit(self, payload: dict[str, Any]) -> None:
"""Submit a task to the stage worker.
@@ -375,22 +366,14 @@ def _stage_worker(
stage_payload: dict[str, Any],
in_q: mp.Queue,
out_q: mp.Queue,
log_file: str | None = None,
batch_timeout: int = 10,
) -> None:
"""Stage worker entry: device setup, LLM init, batching, SHM IPC."""
import logging as _logging
import os as _os
import time as _time
from vllm_omni.distributed.omni_connectors import build_stage_connectors
from vllm_omni.distributed.omni_connectors.adapter import try_recv_via_connector
from vllm_omni.entrypoints.log_utils import (
compute_and_log_stage_request_stats,
count_tokens_from_outputs,
log_stage_batch_stats,
log_stage_running_avg,
)
from vllm_omni.entrypoints.omni_llm import OmniStageLLM
# no inline JSONL/serialization imports; logging handled by utilities
@@ -400,77 +383,6 @@ def _stage_worker(
shm_threshold_bytes = int(stage_payload.get("shm_threshold_bytes", 65536))
connectors_config = stage_payload.get("connectors_config", {})
# Per-stage logger: clear inherited handlers to avoid broken parent streams
try:
stage_log = _logging.getLogger(__name__)
stage_log.setLevel(_logging.DEBUG)
for _h in list(stage_log.handlers):
try:
stage_log.removeHandler(_h)
except Exception:
pass
stage_log.propagate = False
class _StageFilter(_logging.Filter):
def filter(self, record: _logging.LogRecord) -> bool:
setattr(record, "stage", stage_id)
return True
# Prefer file logging even when log_file is not provided, to avoid invalid stdio in child procs
try:
if log_file:
_path = f"{log_file}.stage{stage_id}.log"
else:
_path = f"/tmp/omni_stage{stage_id}.log"
# Ensure dir exists
_os.makedirs(_os.path.dirname(_path), exist_ok=True)
fh = _logging.FileHandler(_path)
fh.setLevel(_logging.DEBUG)
fh.setFormatter(
_logging.Formatter("%(asctime)s [PID:%(process)d] [Stage-%(stage)s] %(levelname)s: %(message)s")
)
class _StageFilter(_logging.Filter):
def filter(self, record: _logging.LogRecord) -> bool:
setattr(record, "stage", stage_id)
return True
fh.addFilter(_StageFilter())
stage_log.addHandler(fh)
try:
# Also route vLLM internal logs to the same handler for this stage
_vllm_logger = _logging.getLogger("vllm")
_vllm_logger.setLevel(_logging.DEBUG)
for _vh in list(_vllm_logger.handlers):
try:
_vllm_logger.removeHandler(_vh)
except Exception:
pass
_vllm_logger.propagate = False
_vllm_logger.addHandler(fh)
except Exception:
pass
except Exception:
# Final fallback: attach NullHandler to avoid logging errors
stage_log.addHandler(_logging.NullHandler())
except Exception:
pass
try:
_logging.raiseExceptions = False
_root_logger = _logging.getLogger()
for _h in list(_root_logger.handlers):
try:
_root_logger.removeHandler(_h)
except Exception:
pass
_root_logger.addHandler(_logging.NullHandler())
except Exception:
pass
# Stage stats JSONL file
_stats_file = f"{log_file}.stage{stage_id}.stats.jsonl" if log_file else None
# Aggregates for running average
_agg_total_tokens = 0
_agg_total_gen_time_ms = 0.0
@@ -485,7 +397,7 @@ def _stage_worker(
device_type = detect_device_type()
set_stage_devices(stage_id, runtime_cfg.get("devices"), device_type=device_type)
except Exception as e:
_logging.getLogger(__name__).warning("[Stage-%s] Device setup failed: %s", stage_id, e)
logger.warning("Device setup failed: %s", e)
# Sequential initialization on the same device to avoid memory calculation errors
# when multiple instances start simultaneously
@@ -537,9 +449,8 @@ def _stage_worker(
# Sort devices_to_lock to prevent deadlock (all processes acquire locks in same order)
devices_to_lock = sorted(devices_to_lock)
_logging.getLogger(__name__).debug(
"[Stage-%s] Parallel config: TP=%d, PP=%d, DP=%d, PCP=%d; will lock %d devices: %s",
stage_id,
logger.debug(
"Parallel config: TP=%d, PP=%d, DP=%d, PCP=%d; will lock %d devices: %s",
tensor_parallel_size,
pipeline_parallel_size,
data_parallel_size,
@@ -572,19 +483,15 @@ def _stage_worker(
_os.fsync(lock_fd) # Ensure written to disk
lock_acquired = True
acquired_lock_fds.append(lock_fd)
_logging.getLogger(__name__).debug(
"[Stage-%s] Acquired exclusive lock for device %s", stage_id, device_id
)
logger.debug("Acquired exclusive lock for device %s", device_id)
except BlockingIOError:
# Lock is held by another process
_os.close(lock_fd)
# Check if we've been waiting too long
if _time.time() - wait_start > max_wait_time:
_logging.getLogger(__name__).warning(
"[Stage-%s] Timeout waiting for device %s "
"initialization lock, proceeding anyway",
stage_id,
logger.warning(
"Timeout waiting for device %s initialization lock, proceeding anyway",
device_id,
)
break
@@ -593,9 +500,8 @@ def _stage_worker(
_time.sleep(0.1)
except OSError as e:
# Other error - log and continue without lock
_logging.getLogger(__name__).debug(
"[Stage-%s] Failed to acquire lock for device %s: %s, continuing anyway",
stage_id,
logger.debug(
"Failed to acquire lock for device %s: %s, continuing anyway",
device_id,
e,
)
@@ -607,14 +513,10 @@ def _stage_worker(
lock_files = acquired_lock_fds
except Exception as e:
_logging.getLogger(__name__).debug(
"[Stage-%s] Failed to set up sequential initialization lock: %s", stage_id, e
)
logger.debug("Failed to set up sequential initialization lock: %s", e)
# Init LLM
_logging.getLogger(__name__).debug(
"[Stage-%s] Initializing engine with args keys=%s", stage_id, list(engine_args.keys())
)
logger.debug("Initializing engine with args keys=%s", list(engine_args.keys()))
try:
stage_engine = OmniStageLLM(model=model, **engine_args)
finally:
@@ -624,10 +526,10 @@ def _stage_worker(
for lock_fd in lock_files:
try:
_os.close(lock_fd)
_logging.getLogger(__name__).debug("[Stage-%s] Released initialization lock (fd=%s)", stage_id, lock_fd)
logger.debug("Released initialization lock (fd=%s)", lock_fd)
except (OSError, ValueError):
pass
_logging.getLogger(__name__).debug("[Stage-%s] Engine initialized", stage_id)
logger.debug("Engine initialized")
# Initialize OmniConnectors if configured
connectors = {}
@@ -646,16 +548,17 @@ def _stage_worker(
except Exception:
pass
max_batch_size = int(runtime_cfg.get("max_batch_size", 1) or 1)
logger.info(f"Max batch size: {max_batch_size}")
# Batch processing loop
while True:
task = in_q.get()
_recv_dequeue_ts = _time.time()
if task is None:
_logging.getLogger(__name__).error("[Stage-%s] Received shutdown signal", stage_id)
logger.error("Received shutdown signal")
break
max_batch_size = int(runtime_cfg.get("max_batch_size", 1) or 1)
print(f"[Stage-{stage_id}] Max batch size: {max_batch_size}")
batch_tasks: list[dict[str, Any]] = [task]
start_time = _time.time()
if max_batch_size > 1:
@@ -711,9 +614,8 @@ def _stage_worker(
"Ensure connectors are configured for all incoming edges."
)
if _rx_metrics:
_rx_decode_ms_by_rid[rid] = float(_rx_metrics.get("rx_decode_time_ms", 0.0))
_rx_bytes_by_rid[rid] = int(_rx_metrics.get("rx_transfer_bytes", 0))
_rx_decode_ms_by_rid[rid] = float(_rx_metrics.get("rx_decode_time_ms", 0.0))
_rx_bytes_by_rid[rid] = int(_rx_metrics.get("rx_transfer_bytes", 0))
batch_request_ids.append(rid)
if isinstance(ein, list):
@@ -721,20 +623,13 @@ def _stage_worker(
elif isinstance(ein, dict):
batch_engine_inputs.append(ein)
else:
_logging.getLogger(__name__).error("[Stage-%s] Invalid engine input type: %s", stage_id, type(ein))
logger.error("Invalid engine input type: %s", type(ein))
sampling_params = batch_tasks[0]["sampling_params"]
_logging.getLogger(__name__).debug(
"[Stage-%s] Received batch size=%d, request_ids=%s",
stage_id,
logger.debug(
"Received batch size=%d, request_ids=%s",
len(batch_tasks),
batch_request_ids,
)
print("--------------------------------", flush=True)
print(
f"[Stage-{stage_id}] Received batch size={len(batch_tasks)}, request_ids={batch_request_ids}",
flush=True,
)
print("--------------------------------", flush=True)
try:
_batch_seq += 1
gen_outputs: list[Any] = []
@@ -743,14 +638,7 @@ def _stage_worker(
gen_outputs.append(ro)
_gen_t1 = _time.time()
_gen_ms = (_gen_t1 - _gen_t0) * 1000.0
try:
print(
f"[Stage-{stage_id}] Generate done: batch={len(batch_tasks)}, "
f"req_ids={batch_request_ids}, gen_ms={_gen_ms:.1f}",
flush=True,
)
except Exception:
pass
logger.debug(f"Generate done: batch={len(batch_tasks)}, req_ids={batch_request_ids}, gen_ms={_gen_ms:.1f}")
# Group outputs per request id with fallback
req_to_outputs: dict[Any, list[Any]] = {rid: [] for rid in batch_request_ids}
@@ -769,56 +657,27 @@ def _stage_worker(
req_to_outputs[target_rid].append(ro)
idx += 1
# Per-request stats logging and aggregates
for rid in batch_request_ids:
_r_outputs = req_to_outputs.get(rid, [])
_num_tokens = count_tokens_from_outputs(_r_outputs)
_agg_total_tokens += _num_tokens
_agg_total_gen_time_ms += _gen_ms
if _stats_file:
_avg_tokens_per_s = (
(_agg_total_tokens * 1000.0 / _agg_total_gen_time_ms) if _agg_total_gen_time_ms > 0 else 0.0
)
log_stage_running_avg(
_stats_file,
stage_id,
int(_agg_total_tokens),
float(_agg_total_gen_time_ms),
float(_avg_tokens_per_s),
)
log_stage_batch_stats(
_stats_file,
stage_id,
len(batch_tasks),
float(_gen_ms),
list(batch_request_ids),
)
_agg_total_gen_time_ms += _gen_ms
# Emit per-request results
for rid in batch_request_ids:
for i, rid in enumerate(batch_request_ids):
r_outputs = req_to_outputs.get(rid, [])
_metrics = make_request_stats(
r_outputs,
_gen_ms,
int(_batch_seq),
int(len(batch_request_ids)),
float(_rx_decode_ms_by_rid.get(rid, 0.0)),
int(_rx_bytes_by_rid.get(rid, 0)),
float(_in_flight_ms_by_rid.get(rid, 0.0)),
)
_agg_total_tokens += _metrics.num_tokens_out
if i == len(batch_request_ids) - 1:
_metrics.stage_stats = make_stage_stats(_agg_total_tokens, _agg_total_gen_time_ms)
else:
_metrics.stage_stats = None
try:
use_shm, payload = maybe_dump_to_shm(r_outputs, shm_threshold_bytes)
_metrics = {
"num_tokens_out": int(count_tokens_from_outputs(r_outputs)),
"stage_gen_time_ms": _gen_ms,
"batch_id": int(_batch_seq),
"rx_decode_time_ms": float(_rx_decode_ms_by_rid.get(rid, 0.0)),
"rx_transfer_bytes": int(_rx_bytes_by_rid.get(rid, 0)),
"rx_in_flight_time_ms": float(_in_flight_ms_by_rid.get(rid, 0.0)),
}
if _stats_file:
compute_and_log_stage_request_stats(
_stats_file,
stage_id,
rid,
len(batch_tasks),
r_outputs,
float(_gen_ms),
int(_metrics["rx_transfer_bytes"]), # type: ignore[index]
float(_metrics["rx_decode_time_ms"]), # type: ignore[index]
)
if use_shm:
out_q.put(
{
@@ -843,22 +702,15 @@ def _stage_worker(
"request_id": rid,
"stage_id": stage_id,
"engine_outputs": r_outputs,
"metrics": {
"num_tokens_out": int(count_tokens_from_outputs(r_outputs)),
"stage_gen_time_ms": _gen_ms,
"rx_decode_time_ms": float(_rx_decode_ms_by_rid.get(rid, 0.0)),
"rx_transfer_bytes": int(_rx_bytes_by_rid.get(rid, 0)),
"rx_in_flight_time_ms": float(_in_flight_ms_by_rid.get(rid, 0.0)),
},
"metrics": _metrics,
}
)
_logging.getLogger(__name__).debug(
"[Stage-%s] Enqueued result for request %s to downstream",
stage_id,
logger.debug(
"Enqueued result for request %s to downstream",
rid,
)
except Exception as e:
_logging.getLogger(__name__).exception("[Stage-%s] Failed on batch %s: %s", stage_id, batch_request_ids, e)
logger.exception("Failed on batch %s: %s", batch_request_ids, e)
_tb = traceback.format_exc()
for rid in batch_request_ids:
out_q.put(
@@ -887,19 +739,12 @@ async def _stage_worker_async(
batch_timeout: int = 10,
) -> None:
"""Stage worker entry: device setup, LLM init, batching, SHM IPC."""
import logging as _logging
import os as _os
import time as _time
from vllm_omni.distributed.omni_connectors import build_stage_connectors
from vllm_omni.distributed.omni_connectors.adapter import try_recv_via_connector
from vllm_omni.entrypoints.async_omni import AsyncOmniStageLLM
from vllm_omni.entrypoints.log_utils import (
compute_and_log_stage_request_stats,
count_tokens_from_outputs,
log_stage_batch_stats,
log_stage_running_avg,
)
# no inline JSONL/serialization imports; logging handled by utilities
@@ -909,32 +754,8 @@ async def _stage_worker_async(
shm_threshold_bytes = int(stage_payload.get("shm_threshold_bytes", 65536))
connectors_config = stage_payload.get("connectors_config", {})
log_file = omni_stage._log_file
in_q = omni_stage._in_q
out_q = omni_stage._out_q
# Per-stage file logger (optional)
try:
if log_file:
stage_log = _logging.getLogger(__name__)
stage_log.setLevel(_logging.DEBUG)
fh = _logging.FileHandler(f"{log_file}.stage{stage_id}.log")
fh.setLevel(_logging.DEBUG)
fh.setFormatter(
_logging.Formatter("%(asctime)s [PID:%(process)d] [Stage-%(stage)s] %(levelname)s: %(message)s")
) # noqa: E501
class _StageFilter(_logging.Filter):
def filter(self, record: _logging.LogRecord) -> bool:
setattr(record, "stage", stage_id)
return True
fh.addFilter(_StageFilter())
stage_log.addHandler(fh)
except Exception:
pass
# Stage stats JSONL file
_stats_file = f"{log_file}.stage{stage_id}.stats.jsonl" if log_file else None
# Aggregates for running average
_agg_total_tokens = 0
@@ -951,7 +772,7 @@ async def _stage_worker_async(
device_type = detect_device_type()
set_stage_devices(stage_id, runtime_cfg.get("devices"), device_type=device_type)
except Exception as e:
_logging.getLogger(__name__).warning("[Stage-%s] Device setup failed: %s", stage_id, e)
logger.warning("Device setup failed: %s", e)
# Initialize OmniConnectors if configured to match sync worker behavior
connectors: dict[Any, Any] = {}
@@ -1014,9 +835,8 @@ async def _stage_worker_async(
# Sort devices_to_lock to prevent deadlock (all processes acquire locks in same order)
devices_to_lock = sorted(devices_to_lock)
_logging.getLogger(__name__).debug(
"[Stage-%s] Parallel config: TP=%d, PP=%d, DP=%d, PCP=%d; will lock %d devices: %s",
stage_id,
logger.debug(
"Parallel config: TP=%d, PP=%d, DP=%d, PCP=%d; will lock %d devices: %s",
tensor_parallel_size,
pipeline_parallel_size,
data_parallel_size,
@@ -1049,19 +869,15 @@ async def _stage_worker_async(
_os.fsync(lock_fd) # Ensure written to disk
lock_acquired = True
acquired_lock_fds.append(lock_fd)
_logging.getLogger(__name__).debug(
"[Stage-%s] Acquired exclusive lock for device %s", stage_id, device_id
)
logger.debug("Acquired exclusive lock for device %s", device_id)
except BlockingIOError:
# Lock is held by another process
_os.close(lock_fd)
# Check if we've been waiting too long
if _time.time() - wait_start > max_wait_time:
_logging.getLogger(__name__).warning(
"[Stage-%s] Timeout waiting for device %s "
"initialization lock, proceeding anyway",
stage_id,
logger.warning(
"Timeout waiting for device %s initialization lock, proceeding anyway",
device_id,
)
break
@@ -1070,9 +886,8 @@ async def _stage_worker_async(
_time.sleep(0.1)
except OSError as e:
# Other error - log and continue without lock
_logging.getLogger(__name__).debug(
"[Stage-%s] Failed to acquire lock for device %s: %s, continuing anyway",
stage_id,
logger.debug(
"Failed to acquire lock for device %s: %s, continuing anyway",
device_id,
e,
)
@@ -1084,14 +899,11 @@ async def _stage_worker_async(
lock_files = acquired_lock_fds
except Exception as e:
_logging.getLogger(__name__).debug(
"[Stage-%s] Failed to set up sequential initialization lock: %s", stage_id, e
)
logger.debug("Failed to set up sequential initialization lock: %s", e)
# Init LLM
_logging.getLogger(__name__).debug(
"[Stage-%s] Initializing engine with args keys=%s",
stage_id,
logger.debug(
"Initializing engine with args keys=%s",
list(engine_args.keys()),
)
try:
@@ -1110,13 +922,13 @@ async def _stage_worker_async(
for lock_fd in lock_files:
try:
_os.close(lock_fd)
_logging.getLogger(__name__).debug("[Stage-%s] Released initialization lock (fd=%s)", stage_id, lock_fd)
logger.debug("Released initialization lock (fd=%s)", lock_fd)
except (OSError, ValueError):
pass
omni_stage.set_async_engine(stage_engine)
# Don't keep the dummy data in memory
await stage_engine.reset_mm_cache()
_logging.getLogger(__name__).debug("[Stage-%s] Engine initialized", stage_id)
logger.debug("Engine initialized")
# Signal readiness to orchestrator and send vllm_config back to main process
try:
# Send vllm_config back to main process so it can be accessed via
@@ -1135,8 +947,9 @@ async def _stage_worker_async(
}
)
except Exception as e:
_logging.getLogger(__name__).warning("[Stage-%s] Failed to send stage ready signal: %s", stage_id, e)
logger.warning("Failed to send stage ready signal: %s", e)
generation_out_q = asyncio.Queue()
# Batch processing loop
_rx_bytes_by_rid: dict[Any, int] = {}
_rx_decode_ms_by_rid: dict[Any, float] = {}
@@ -1168,10 +981,7 @@ async def _stage_worker_async(
_rx_bytes_by_rid[rid] = int(_rx_metrics.get("rx_transfer_bytes", 0))
sampling_params = task["sampling_params"]
_logging.getLogger(__name__).debug("[Stage-%s] Received batch size=1, request_ids=%s", stage_id, rid)
print("--------------------------------", flush=True)
print(f"[Stage-{stage_id}] Received batch size=1, request_ids={rid}", flush=True)
print("--------------------------------", flush=True)
logger.debug("Received batch size=1, request_ids=%s", rid)
_gen_t0 = _time.time()
if isinstance(ein, list):
ein = ein[0]
@@ -1181,7 +991,7 @@ async def _stage_worker_async(
_gen_ms = (_gen_t1 - _gen_t0) * 1000.0
await generation_out_q.put((rid, gen_output, _gen_ms))
except Exception as e:
_logging.getLogger(__name__).exception("[Stage-%s] Failed on request %s: %s", stage_id, rid, e)
logger.exception("Failed on request %s: %s", rid, e)
out_q.put(
{
"request_id": rid,
@@ -1195,7 +1005,7 @@ async def _stage_worker_async(
try:
task = in_q.get_nowait()
if task is None:
_logging.getLogger(__name__).debug("[Stage-%s] Received shutdown signal", stage_id)
logger.debug("Received shutdown signal")
break
asyncio.create_task(generation_single_request(task))
except queue.Empty:
@@ -1203,14 +1013,24 @@ async def _stage_worker_async(
batch_request_outputs: list[Any] = []
batch_request_ids: list[Any] = []
_gen_ms_list = []
batch_metrics: list[Any] = []
while True:
try:
rids, gen_output, _gen_ms = generation_out_q.get_nowait()
_num_tokens = count_tokens_from_outputs([gen_output])
rid, gen_output, _gen_ms = generation_out_q.get_nowait()
_metrics = make_request_stats(
[gen_output],
_gen_ms,
int(_batch_seq),
1, # temporarily set to 1
float(_rx_decode_ms_by_rid.get(rid, 0.0)),
int(_rx_bytes_by_rid.get(rid, 0)),
float(_in_flight_ms_by_rid.get(rid, 0.0)),
)
batch_metrics.append(_metrics)
batch_request_outputs.append(gen_output)
_gen_ms_list.append(_gen_ms)
batch_request_ids.append(rids)
_agg_total_tokens += _num_tokens
batch_request_ids.append(rid)
_agg_total_tokens += _metrics.num_tokens_out
except asyncio.QueueEmpty:
await asyncio.sleep(0.001)
break
@@ -1218,48 +1038,22 @@ async def _stage_worker_async(
if not batch_request_outputs:
continue
_batch_seq += 1
if _stats_file:
_batch_gen_t1 = _time.time()
_agg_total_gen_time_ms += (_batch_gen_t1 - _batch_gen_t0) * 1000
_batch_gen_t0 = _batch_gen_t1
_avg_tokens_per_s = (
(_agg_total_tokens * 1000.0 / _agg_total_gen_time_ms) if _agg_total_gen_time_ms > 0 else 0.0
)
log_stage_running_avg(
_stats_file,
stage_id,
int(_agg_total_tokens),
float(_agg_total_gen_time_ms),
float(_avg_tokens_per_s),
)
logger.info("[Stage-%s] Running avg: %s tokens/s", stage_id, _avg_tokens_per_s)
for rid, _gen_ms in zip(batch_request_ids, _gen_ms_list):
log_stage_batch_stats(_stats_file, stage_id, 1, float(_gen_ms), [rid])
logger.info("[Stage-%s] Sending outputs to main process", stage_id)
for rid, output, _gen_ms in zip(batch_request_ids, batch_request_outputs, _gen_ms_list):
_batch_gen_t1 = _time.time()
_agg_total_gen_time_ms += (_batch_gen_t1 - _batch_gen_t0) * 1000
_batch_gen_t0 = _batch_gen_t1
for idx, metrics in enumerate(batch_metrics):
metrics.batch_size = len(batch_metrics)
if idx == len(batch_metrics) - 1:
metrics.stage_stats = make_stage_stats(_agg_total_tokens, _agg_total_gen_time_ms)
logger.debug("Sending outputs to main process")
for rid, output, _gen_ms, _metrics in zip(
batch_request_ids, batch_request_outputs, _gen_ms_list, batch_metrics
):
try:
r_outputs = [output]
use_shm, payload = maybe_dump_to_shm(r_outputs, shm_threshold_bytes)
_metrics = {
"num_tokens_out": int(count_tokens_from_outputs(r_outputs)),
"stage_gen_time_ms": _gen_ms,
"batch_id": int(_batch_seq),
"rx_decode_time_ms": float(_rx_decode_ms_by_rid.pop(rid, 0.0)),
"rx_transfer_bytes": int(_rx_bytes_by_rid.pop(rid, 0)),
"rx_in_flight_time_ms": float(_in_flight_ms_by_rid.pop(rid, 0.0)),
}
if _stats_file:
compute_and_log_stage_request_stats(
_stats_file,
stage_id,
rid,
1,
r_outputs,
float(_gen_ms),
int(_metrics["rx_transfer_bytes"]), # type: ignore[index]
float(_metrics["rx_decode_time_ms"]), # type: ignore[index]
)
if use_shm:
out_q.put(
{
@@ -1278,18 +1072,10 @@ async def _stage_worker_async(
"metrics": _metrics,
}
)
try:
print(
f"[Stage-{stage_id}] Enqueued req={rid}, use_shm={use_shm}, "
f"tokens_out={_metrics['num_tokens_out']}",
flush=True,
)
except Exception:
pass
logger.debug(f"Enqueued req={rid}, use_shm={use_shm}, tokens_out={_metrics.num_tokens_out}")
except Exception as e:
_logging.getLogger(__name__).exception(
"[Stage-%s] Failed to enqueue result for request %s: %s",
stage_id,
logger.exception(
"Failed to enqueue result for request %s: %s",
rid,
e,
)
@@ -1298,17 +1084,42 @@ async def _stage_worker_async(
"request_id": rid,
"stage_id": stage_id,
"engine_outputs": r_outputs,
"metrics": {
"num_tokens_out": int(count_tokens_from_outputs(r_outputs)),
"stage_gen_time_ms": _gen_ms,
"rx_decode_time_ms": float(_rx_decode_ms_by_rid.pop(rid, 0.0)),
"rx_transfer_bytes": int(_rx_bytes_by_rid.pop(rid, 0)),
"rx_in_flight_time_ms": float(_in_flight_ms_by_rid.pop(rid, 0.0)),
},
"metrics": _metrics,
}
)
_logging.getLogger(__name__).debug("[Stage-%s] Enqueued result for request %s to downstream", stage_id, rid)
logger.debug("Enqueued result for request %s to downstream", rid)
print("--------------------------------", flush=True)
print(f"[Stage-{stage_id}] Stage worker exiting", flush=True)
print("--------------------------------", flush=True)
logger.info("Stage worker exiting")
def make_request_stats(
req_output: list[Any],
stage_gen_time_ms: float,
batch_id: int,
batch_size: int,
rx_decode_time_ms: float,
rx_transfer_bytes: int,
rx_in_flight_time_ms: float,
):
from vllm_omni.entrypoints.log_utils import (
StageRequestMetrics,
count_tokens_from_outputs,
)
num_tokens_out = count_tokens_from_outputs(req_output)
return StageRequestMetrics(
num_tokens_out=num_tokens_out,
stage_gen_time_ms=stage_gen_time_ms,
batch_id=batch_id,
batch_size=batch_size,
rx_decode_time_ms=rx_decode_time_ms,
rx_transfer_bytes=rx_transfer_bytes,
rx_in_flight_time_ms=rx_in_flight_time_ms,
stage_stats=None,
)
def make_stage_stats(_agg_total_tokens: int, _agg_total_gen_time_ms: float):
from vllm_omni.entrypoints.log_utils import StageStats
return StageStats(total_token=_agg_total_tokens, total_gen_time=_agg_total_gen_time_ms)