Add pid prefix to process stdout/stderr instead of logger

Author: njhill

vllm/engine/local_worker_utils.py

@@ -1,10 +1,12 @@
 import asyncio
 import multiprocessing
 import os
+import sys
 import threading
 import traceback
 import uuid
 from dataclasses import dataclass
+from io import TextIOBase
 from multiprocessing.connection import wait
 from typing import Any, Callable, Dict, Generic, List, Optional, TypeVar, Union
 
@@ -16,6 +18,10 @@
 
 _TERMINATE = "TERMINATE"  # sentinel
 
+# ANSI color codes
+CYAN = '\033[1;36m'
+RESET = '\033[0;0m'
+
 # Use dedicated multiprocess context for workers.
 # Both spawn and fork work
 mp_method = os.getenv("MULTIPROC_METHOD", "fork")
@@ -172,9 +178,11 @@ def kill_worker(self):
         self.kill()
 
     def run(self) -> None:
-        # Re-init logger in forked process, to include worker-specific prefix
-        global logger
-        logger = init_logger(__name__)
+        # Add process-specific prefix to stdout and stderr
+        process_name = mp.current_process().name
+        pid = os.getpid()
+        _add_prefix(sys.stdout, process_name, pid)
+        _add_prefix(sys.stderr, process_name, pid)
 
         del self.tasks  # Not used in forked process
         self.worker = self.worker_factory()
@@ -205,3 +213,30 @@ def run(self) -> None:
             logger.exception("Worker failed")
 
         logger.info("Worker exiting")
+
+
+def _add_prefix(file: TextIOBase, worker_name: str, pid: int) -> None:
+    """Prepend output with process-specific prefix"""
+
+    prefix = f"{CYAN}({worker_name} pid={pid}){RESET} "
+    file_write = file.write
+
+    def write_with_prefix(s: str):
+        if not s:
+            return
+        if file.start_new_line:
+            file_write(prefix)
+        idx = 0
+        while (next_idx := s.find('\n', idx)) != -1:
+            next_idx += 1
+            file_write(s[idx:next_idx])
+            if next_idx == len(s):
+                file.start_new_line = True
+                return
+            file_write(prefix)
+            idx = next_idx
+        file_write(s[idx:])
+        file.start_new_line = False
+
+    file.start_new_line = True
+    file.write = write_with_prefix

vllm/logger.py

@@ -2,7 +2,6 @@
 # https://github.com/skypilot-org/skypilot/blob/86dc0f6283a335e4aa37b3c10716f90999f48ab6/sky/sky_logging.py
 """Logging configuration for vLLM."""
 import logging
-import multiprocessing as mp
 import os
 import sys
 from typing import Optional
@@ -13,13 +12,6 @@
 _DATE_FORMAT = "%m-%d %H:%M:%S"
 
 
-class ProcessLogger(logging.LoggerAdapter):
-
-    def process(self, msg, kwargs):
-        msg = f"[{self.extra['process_name']} pid {self.extra['pid']}] {msg}"
-        return msg, kwargs
-
-
 class NewLineFormatter(logging.Formatter):
     """Adds logging prefix to newlines to align multi-line messages."""
 
@@ -72,9 +64,4 @@ def init_logger(name: str):
                 " Please open an issue on Github.")
         logger.addHandler(_default_handler)
         logger.propagate = False
-    if mp.parent_process() is not None:
-        logger = ProcessLogger(logger, {
-            'process_name': mp.current_process().name,
-            'pid': os.getpid()
-        })
     return logger