From a500ca1f28f191863cfef7698cc889ca9941b635 Mon Sep 17 00:00:00 2001 From: dengyunyang <584797741@qq.com> Date: Fri, 26 Dec 2025 08:42:06 +0800 Subject: [PATCH] [Feature] logger refactor (#466) Signed-off-by: dengyunyang <584797741@qq.com> Co-authored-by: Hongsheng Liu --- tests/entrypoints/test_omni_llm.py | 2 +- vllm_omni/entrypoints/async_omni.py | 24 +- vllm_omni/entrypoints/log_utils.py | 356 ++++++------------ vllm_omni/entrypoints/omni_llm.py | 24 +- vllm_omni/entrypoints/omni_stage.py | 549 +++++++++------------------- 5 files changed, 310 insertions(+), 645 deletions(-) diff --git a/tests/entrypoints/test_omni_llm.py b/tests/entrypoints/test_omni_llm.py index 1b36ed509..cef5dce43 100644 --- a/tests/entrypoints/test_omni_llm.py +++ b/tests/entrypoints/test_omni_llm.py @@ -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 diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index f57a90465..e59b02778 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -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.log" - ) logger.error( "[Orchestrator] Stage initialization failed, shutting down. \ Suggestions:\n- %s", diff --git a/vllm_omni/entrypoints/log_utils.py b/vllm_omni/entrypoints/log_utils.py index 241308e07..ac5bfc277 100644 --- a/vllm_omni/entrypoints/log_utils.py +++ b/vllm_omni/entrypoints/log_utils.py @@ -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 diff --git a/vllm_omni/entrypoints/omni_llm.py b/vllm_omni/entrypoints/omni_llm.py index bfca19573..36cd60a5a 100644 --- a/vllm_omni/entrypoints/omni_llm.py +++ b/vllm_omni/entrypoints/omni_llm.py @@ -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.log") logger.error( "[Orchestrator] Stage initialization failed, shutting down. Suggestions:\n- %s", "\n- ".join(suggestions), diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 503c8395c..692c122f8 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -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)