utils.py

import os
import weakref
from collections.abc import Sequence
from typing import (Any, Callable, Dict, Generic, List, Optional, TypeVar,
                    Union, overload)

from vllm.logger import init_logger
from vllm.utils import get_mp_context, kill_process_tree

logger = init_logger(__name__)

T = TypeVar("T")


class ConstantList(Generic[T], Sequence):

    def __init__(self, x: List[T]) -> None:
        self._x = x

    def append(self, item):
        raise Exception("Cannot append to a constant list")

    def extend(self, item):
        raise Exception("Cannot extend a constant list")

    def insert(self, item):
        raise Exception("Cannot insert into a constant list")

    def pop(self, item):
        raise Exception("Cannot pop from a constant list")

    def remove(self, item):
        raise Exception("Cannot remove from a constant list")

    def clear(self):
        raise Exception("Cannot clear a constant list")

    def index(self,
              item: T,
              start: int = 0,
              stop: Optional[int] = None) -> int:
        return self._x.index(item, start,
                             stop if stop is not None else len(self._x))

    @overload
    def __getitem__(self, item: int) -> T:
        ...

    @overload
    def __getitem__(self, s: slice, /) -> List[T]:
        ...

    def __getitem__(self, item: Union[int, slice]) -> Union[T, List[T]]:
        return self._x[item]

    @overload
    def __setitem__(self, item: int, value: T):
        ...

    @overload
    def __setitem__(self, s: slice, value: T, /):
        ...

    def __setitem__(self, item: Union[int, slice], value: Union[T, List[T]]):
        raise Exception("Cannot set item in a constant list")

    def __delitem__(self, item):
        raise Exception("Cannot delete item from a constant list")

    def __iter__(self):
        return iter(self._x)

    def __contains__(self, item):
        return item in self._x

    def __len__(self):
        return len(self._x)


class BackgroundProcHandle:
    """
    Utility class to handle creation, readiness, and shutdown
    of background processes used by the AsyncLLM and LLMEngine.
    """

    def __init__(
        self,
        input_path: str,
        output_path: str,
        process_name: str,
        target_fn: Callable,
        process_kwargs: Dict[Any, Any],
    ):
        self._finalizer = weakref.finalize(self, self.shutdown)

        context = get_mp_context()
        reader, writer = context.Pipe(duplex=False)

        assert ("ready_pipe" not in process_kwargs
                and "input_path" not in process_kwargs
                and "output_path" not in process_kwargs)
        process_kwargs["ready_pipe"] = writer
        process_kwargs["input_path"] = input_path
        process_kwargs["output_path"] = output_path
        self.input_path = input_path
        self.output_path = output_path

        # Run Detokenizer busy loop in background process.
        self.proc = context.Process(target=target_fn, kwargs=process_kwargs)
        self.proc.start()

        # Wait for startup.
        if reader.recv()["status"] != "READY":
            raise RuntimeError(f"{process_name} initialization failed. "
                               "See root cause above.")

    def __del__(self):
        self.shutdown()

    def shutdown(self):
        # Shutdown the process if needed.
        if hasattr(self, "proc") and self.proc.is_alive():
            self.proc.terminate()
            self.proc.join(5)

            if self.proc.is_alive():
                kill_process_tree(self.proc.pid)

        # Remove zmq ipc socket files
        ipc_sockets = [self.output_path, self.input_path]
        for ipc_socket in ipc_sockets:
            socket_file = ipc_socket.replace("ipc://", "")
            if os and os.path.exists(socket_file):
                os.remove(socket_file)