vllm.v1.utils ¶

class APIServerProcessManager:
    """Manages a group of API server processes.

    Handles creation, monitoring, and termination of API server worker
    processes. Also monitors extra processes to check if they are healthy.
    """

    def __init__(
        self,
        listen_address: str,
        sock: Any,
        args: argparse.Namespace,
        num_servers: int,
        input_addresses: list[str],
        output_addresses: list[str],
        target_server_fn: Callable | None = None,
        stats_update_address: str | None = None,
        tensor_queue: Queue | None = None,
    ):
        """Initialize and start API server worker processes.

        ``input_addresses``/``output_addresses`` may contain
        ``tcp://host:0`` placeholders; each child must report the actual
        bound endpoint over its ``actual_address_pipe`` in ``client_config``
        and the parent collects them via
        :py:meth:`gather_actual_addresses`.

        Args:
            target_server_fn: Override function to call for each API server process
            listen_address: Address to listen for client connections
            sock: Socket for client connections
            args: Command line arguments
            num_servers: Number of API server processes to start
            input_addresses: Input addresses for each API server
            output_addresses: Output addresses for each API server
            stats_update_address: Optional stats update address
            tensor_queue: Optional tensor IPC queue for sharing MM tensors
        """
        self.listen_address = listen_address
        self.sock = sock
        self.args = args

        spawn_context = multiprocessing.get_context("spawn")
        self.processes: list[BaseProcess] = []
        self._address_pipes: list[connection.Connection] = []

        for i, in_addr, out_addr in zip(
            range(num_servers), input_addresses, output_addresses
        ):
            client_config: dict[str, Any] = {
                "input_address": in_addr,
                "output_address": out_addr,
                "client_count": num_servers,
                "client_index": i,
            }
            if stats_update_address is not None:
                client_config["stats_update_address"] = stats_update_address
            if tensor_queue is not None:
                client_config["tensor_queue"] = tensor_queue

            parent_recv, child_send = spawn_context.Pipe(duplex=False)
            self._address_pipes.append(parent_recv)
            client_config["actual_address_pipe"] = child_send

            proc = spawn_context.Process(
                target=target_server_fn or run_api_server_worker_proc,
                name=f"ApiServer_{i}",
                args=(listen_address, sock, args, client_config),
            )
            self.processes.append(proc)
            proc.start()

            # Drop parent's write end so reader sees EOF on child death.
            child_send.close()

        logger.info("Started %d API server processes", len(self.processes))

        # Shutdown only the API server processes on garbage collection
        # The extra processes are managed by their owners
        self._finalizer = weakref.finalize(self, shutdown, self.processes)

    def gather_actual_addresses(
        self,
        timeout: float = envs.VLLM_ENGINE_READY_TIMEOUT_S,
    ) -> tuple[list[str], list[str]]:
        """Return (inputs, outputs) reported by each child, indexed by
        ``client_index``. Raises ``RuntimeError`` on timeout or premature
        child exit."""
        n = len(self._address_pipes)
        inputs: list[str | None] = [None] * n
        outputs: list[str | None] = [None] * n
        pending: dict[connection.Connection, int] = {
            pipe: i for i, pipe in enumerate(self._address_pipes)
        }
        sentinel_to_idx: dict[Any, int] = {
            proc.sentinel: i for i, proc in enumerate(self.processes)
        }

        deadline = time.monotonic() + timeout
        try:
            while pending:
                remaining = deadline - time.monotonic()
                if remaining <= 0:
                    missing = [self.processes[i].name for i in pending.values()]
                    raise RuntimeError(
                        f"Timed out after {timeout:.1f}s waiting for "
                        f"API server(s) to report bound ZMQ addresses: "
                        f"{missing}"
                    )
                waitables: list[Any] = list(pending.keys()) + list(
                    sentinel_to_idx.keys()
                )
                ready = connection.wait(waitables, timeout=remaining)
                # Drain pipes before checking sentinels: a child that sent
                # its message and then exited can surface both events in
                # the same poll, and we must record the success first.
                for item in ready:
                    if isinstance(item, connection.Connection) and item in pending:
                        idx = pending.pop(item)
                        try:
                            msg: dict[str, str] = item.recv()
                        except EOFError as e:
                            raise RuntimeError(
                                f"API server {self.processes[idx].name} "
                                f"closed its address pipe without "
                                f"reporting its bound ZMQ addresses"
                            ) from e
                        inputs[idx] = msg["input_address"]
                        outputs[idx] = msg["output_address"]
                        item.close()
                for item in ready:
                    if item in sentinel_to_idx:
                        idx = sentinel_to_idx.pop(item)
                        pipe = self._address_pipes[idx]
                        if pipe in pending:
                            proc = self.processes[idx]
                            raise RuntimeError(
                                f"API server process {proc.name} exited "
                                f"(code={proc.exitcode}) before reporting "
                                f"its bound ZMQ addresses"
                            )
        finally:
            for pipe in pending:
                with contextlib.suppress(Exception):
                    pipe.close()

        return inputs, outputs  # type: ignore[return-value]

    def shutdown(self, timeout: float | None = None) -> None:
        """Shutdown API server processes with configurable timeout"""
        for pipe in self._address_pipes:
            with contextlib.suppress(Exception):
                pipe.close()
        self._address_pipes = []

        if self._finalizer.detach() is not None:
            shutdown(self.processes, timeout=timeout)

class CpuGpuBuffer:
    """Buffer to easily copy tensors between CPU and GPU."""

    def __init__(
        self,
        *size: int | torch.SymInt,
        dtype: torch.dtype,
        device: torch.device,
        pin_memory: bool,
        with_numpy: bool = True,
    ) -> None:
        # these buffers are mutable runtime state, so allocate them as normal
        with torch.inference_mode(False):
            self.cpu = torch.zeros(
                *size, dtype=dtype, device="cpu", pin_memory=pin_memory
            )
            self.gpu = torch.zeros_like(self.cpu, device=device)
        self.np: np.ndarray
        # To keep type hints simple (avoiding generics and subclasses), we
        # only conditionally create the numpy array attribute. This can cause
        # AttributeError if `self.np` is accessed when `with_numpy=False`.
        if with_numpy:
            if dtype == torch.bfloat16:
                raise ValueError(
                    "Bfloat16 torch tensors cannot be directly cast to a "
                    "numpy array, so call CpuGpuBuffer with with_numpy=False"
                )
            self.np = self.cpu.numpy()

    def copy_to_gpu(self, n: int | None = None) -> torch.Tensor:
        if n is None:
            return self.gpu.copy_(self.cpu, non_blocking=True)
        return self.gpu[:n].copy_(self.cpu[:n], non_blocking=True)

    def copy_to_cpu(self, n: int | None = None) -> torch.Tensor:
        """NOTE: Because this method is non-blocking, explicit synchronization
        is needed to ensure the data is copied to CPU."""
        if n is None:
            return self.cpu.copy_(self.gpu, non_blocking=True)
        return self.cpu[:n].copy_(self.gpu[:n], non_blocking=True)

class RustFrontendProcessManager:
    """Manages a single Rust frontend subprocess.

    Launches the Rust vllm-rs binary in 'frontend' mode, passing the
    listening socket fd and ZMQ transport addresses. Provides the same
    interface as APIServerProcessManager for process monitoring.
    """

    def __init__(
        self,
        binary_path: str,
        sock: Any,
        args: argparse.Namespace,
        input_address: str,
        output_address: str,
        engine_count: int,
        stats_update_address: str | None = None,
    ):
        import os
        import subprocess

        fd = sock.fileno()
        os.set_inheritable(fd, True)

        cmd = [
            binary_path,
            "frontend",
            "--listen-fd",
            str(fd),
            "--input-address",
            input_address,
            "--output-address",
            output_address,
            "--engine-count",
            str(engine_count),
        ]
        if stats_update_address is not None:
            cmd.extend(["--coordinator-address", stats_update_address])
        from vllm.entrypoints.utils import jsonify_non_default_args

        args_json = json.dumps(
            jsonify_non_default_args(args, exclude={"api_server_count"}),
            sort_keys=True,
        )
        cmd.extend(["--args-json", args_json])

        logger.info("Launching Rust frontend: %s", " ".join(cmd))
        self._proc = subprocess.Popen(cmd, pass_fds=(fd,))

        # Create a process wrapper with a sentinel fd for monitoring
        self.processes: list[_SubprocessWrapper] = [
            _SubprocessWrapper(self._proc, "RustFrontend")
        ]

        self._finalizer = weakref.finalize(self, _shutdown_subprocesses, self.processes)

    def shutdown(self, timeout: float | None = None) -> None:
        if self._finalizer.detach() is not None:
            _shutdown_subprocesses(self.processes, timeout=timeout)

class _SubprocessWrapper:
    """Wraps subprocess.Popen to provide the BaseProcess-like interface
    needed by wait_for_completion_or_failure."""

    def __init__(self, proc, name: str):
        self._proc = proc
        self.name = name
        self.pid = proc.pid
        self._sentinel_conn: connection.Connection | None = None
        self._sentinel_send: connection.Connection | None = None

        # Use a Pipe-based sentinel so subprocess monitoring works uniformly
        # across platforms with multiprocessing.connection.wait().
        recv, send = connection.Pipe(duplex=False)
        self._sentinel_conn = recv
        self._sentinel_send = send

        def monitor_subprocess() -> None:
            try:
                proc.wait()
            finally:
                with contextlib.suppress(Exception):
                    send.close()

        threading.Thread(
            target=monitor_subprocess, daemon=True, name=f"{name}Monitor"
        ).start()

    @property
    def sentinel(self):
        return self._sentinel_conn

    @property
    def exitcode(self) -> int | None:
        return self._proc.returncode if self._proc.poll() is not None else None

    def is_alive(self) -> bool:
        return self._proc.poll() is None

    def terminate(self):
        self._proc.terminate()

    def join(self, timeout=None):
        with contextlib.suppress(Exception):
            self._proc.wait(timeout=timeout)

    def __del__(self):
        with contextlib.suppress(Exception):
            if self._sentinel_conn is not None:
                self._sentinel_conn.close()
            if self._sentinel_send is not None:
                self._sentinel_send.close()