vllm-omni/vllm_omni/edge/local/external/client.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""The engine's half of an out-of-process stage.

Binds a listener, starts the worker, waits for it to connect *back*, and then
speaks :mod:`~vllm_omni.edge.local.external.protocol` over that socket. The
connect-back direction is what keeps this deployable: WSL reaching into Windows
needs an inbound firewall rule, while Windows reaching into WSL does not.

Two things here are load-bearing rather than defensive.

**The placement gate.** :meth:`ExternalWorker.load` returns a
:class:`LoadReport` carrying the per-provider node counts ORT wrote while the
graph actually ran. ``fraction_on_target`` is the number the planner refuses
on. It has to be measured this way because every cheaper check is wrong: a
VitisAI session lists the EP in ``get_providers()`` whether or not it took any
nodes, and when it takes none the outputs are bit-identical to the CPU's -- so
neither the provider list nor the numbers can tell a real NPU run from a silent
fallback. Only the node assignment can.

**The worker's own memory.** The engine's sampler walks a psutil process tree,
which stops at the WSL boundary; a Windows worker's resident set is invisible
to it. Since both AMD devices allocate out of the *same host RAM* the CPU
stage is budgeted against, an unreported worker is a hole in the admission
ledger, not a rounding error. So the worker reports its own RSS and
:meth:`stats` is how that reaches the ledger.
"""

from __future__ import annotations

import secrets
import socket
import subprocess
import sys
import time
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any

import numpy as np

from vllm_omni.edge.local.external import launch as _launch
from vllm_omni.edge.local.external import protocol as _proto
from vllm_omni.edge.local.external.protocol import ProtocolError, WorkerError

DEFAULT_START_TIMEOUT_S = 120.0
"""Generous on purpose: the first VitisAI session on a cold machine spends
over a second in ``InferenceSession`` alone, and a Windows interpreter starting
off a ``\\\\wsl.localhost`` share is slower still."""

DEFAULT_CALL_TIMEOUT_S = 300.0

DEFAULT_CLOSE_TIMEOUT_S = 10.0
"""Bounded, for the same reason M0 bounds its cancellation drain: a worker that
will not exit must not hold the engine's shutdown open."""


@dataclass(frozen=True)
class LoadReport:
    """What the worker did when it opened the graph, and where it put it."""

    ep: str
    graph_path: str
    session_providers: tuple[str, ...] = ()
    available_providers: tuple[str, ...] = ()
    node_counts: dict[str, int] = field(default_factory=dict)
    total_nodes: int = 0
    target_nodes: int = 0
    fraction_on_target: float | None = None
    """``None`` means *unverified*, which is not the same as zero and must not
    be rounded into it: no run happened, so ORT wrote no assignments."""
    session_create_s: float = 0.0
    warmup_s: float = 0.0
    rss_bytes: int = 0
    inputs: tuple[dict[str, Any], ...] = ()
    outputs: tuple[dict[str, Any], ...] = ()
    profile_path: str = ""
    ep_directory: str = ""
    requested_provider_missing: str = ""
    fallback_note: str = ""
    onnxruntime: str = ""
    note: str = ""

    @classmethod
    def from_body(cls, body: dict[str, Any], *, ep: str, graph_path: str) -> LoadReport:
        return cls(
            ep=ep,
            graph_path=graph_path,
            session_providers=tuple(body.get("session_providers") or ()),
            available_providers=tuple(body.get("available_providers") or ()),
            node_counts=dict(body.get("node_counts") or {}),
            total_nodes=int(body.get("total_nodes") or 0),
            target_nodes=int(body.get("target_nodes") or 0),
            fraction_on_target=body.get("fraction_on_target"),
            session_create_s=float(body.get("session_create_s") or 0.0),
            warmup_s=float(body.get("warmup_s") or 0.0),
            rss_bytes=int(body.get("rss_bytes") or 0),
            inputs=tuple(body.get("inputs") or ()),
            outputs=tuple(body.get("outputs") or ()),
            profile_path=str(body.get("profile_path") or ""),
            ep_directory=str(body.get("ep_directory") or ""),
            requested_provider_missing=str(body.get("requested_provider_missing") or ""),
            fallback_note=str(body.get("fallback_note") or ""),
            onnxruntime=str(body.get("onnxruntime") or ""),
            note=str(body.get("placement_unknown_reason") or ""),
        )

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)

    def summary(self) -> str:
        if self.fraction_on_target is None:
            placement = f"placement unverified ({self.note or 'no profiled run'})"
        else:
            split = ", ".join(f"{k}={v}" for k, v in sorted(self.node_counts.items()))
            placement = f"{self.target_nodes}/{self.total_nodes} nodes on {self.ep} ({split})"
        return (
            f"{Path(self.graph_path).name} on {self.ep}: {placement}; "
            f"session {self.session_create_s:.3f}s, warmup {self.warmup_s:.3f}s, "
            f"worker rss {self.rss_bytes / 2**20:.0f} MiB"
        )


@dataclass(frozen=True)
class RunTiming:
    """Where one forward's wall time went.

    Split because the composition rule needs it split: a stage only earns its
    place on a second device when the end-to-end gain beats the cost of getting
    there, and ``transport_s`` is that cost measured rather than assumed.
    """

    worker_s: float
    """Inside ``session.run`` on the far side."""
    round_trip_s: float
    """Wall time here, from send to the last byte of the reply."""

    @property
    def transport_s(self) -> float:
        """Serialisation, both socket traversals, and the worker's dispatch."""
        return max(0.0, self.round_trip_s - self.worker_s)

    def to_dict(self) -> dict[str, float]:
        return {
            "worker_s": self.worker_s,
            "round_trip_s": self.round_trip_s,
            "transport_s": self.transport_s,
        }


class WorkerStartError(RuntimeError):
    """The worker never connected back. Carries whatever it printed."""


def _default_gateway() -> str | None:
    """The far side of the default route, i.e. the Windows host under WSL2 NAT."""
    try:
        lines = Path("/proc/net/route").read_text().splitlines()[1:]
    except OSError:
        return None
    for line in lines:
        fields = line.split()
        if len(fields) > 2 and fields[1] == "00000000":
            packed = int(fields[2], 16)
            return ".".join(str((packed >> shift) & 0xFF) for shift in (0, 8, 16, 24))
    return None


def _bind_address(*, is_windows: bool) -> str:
    """The address the worker should dial.

    A WSL worker gets loopback, which nothing outside this machine can reach. A
    native-Windows worker cannot use loopback -- under WSL2's NAT the two live
    on different network stacks -- so it needs this distro's address on the
    host-only network (``172.21.141.90`` here), which is the source address the
    routing table picks for the *default gateway*.

    The gateway specifically, not any address that happens to be handy. The
    obvious candidate is the nameserver in ``/etc/resolv.conf``, and it is
    wrong: WSL puts ``10.255.255.254`` on a loopback alias, so a probe aimed at
    it reports ``10.255.255.254`` as the source, the listener binds an address
    only this distro can reach, and the Windows worker gets ECONNREFUSED after
    a silent start. That is exactly the failure this comment exists to prevent
    a future edit from reintroducing.
    """
    if not is_windows or _launch.is_wsl() is False:
        return "127.0.0.1"
    # 1.1.1.1 is a fallback route target, not a destination: UDP ``connect``
    # sends no packets, it only resolves which local address would be used.
    for target in (_default_gateway(), "1.1.1.1"):
        if not target:
            continue
        probe = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        try:
            probe.connect((target, 53))
            address = str(probe.getsockname()[0])
        except OSError:
            continue
        finally:
            probe.close()
        if not address.startswith("127."):
            return address
    return "127.0.0.1"


class ExternalWorker:
    """A stage running in another interpreter, possibly on another OS."""

    def __init__(
        self,
        route: _launch.Route,
        *,
        start_timeout_s: float = DEFAULT_START_TIMEOUT_S,
        call_timeout_s: float = DEFAULT_CALL_TIMEOUT_S,
    ) -> None:
        if not route.available:
            raise WorkerStartError(f"route {route.name} is not available: {route.reason}")
        self.route = route
        self.start_timeout_s = start_timeout_s
        self.call_timeout_s = call_timeout_s
        self._process: subprocess.Popen[bytes] | None = None
        self._sock: socket.socket | None = None
        self._listener: socket.socket | None = None
        self.hello: dict[str, Any] = {}
        self.report: LoadReport | None = None

    # -- lifecycle ----------------------------------------------------------

    def start(self) -> dict[str, Any]:
        """Launch the worker and complete the handshake. Returns its ``hello``."""
        token = secrets.token_hex(16)
        host = _bind_address(is_windows=self.route.is_windows)
        listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        listener.bind((host, 0))
        listener.listen(1)
        listener.settimeout(self.start_timeout_s)
        self._listener = listener
        port = listener.getsockname()[1]

        argv = [
            str(self.route.interpreter),
            _launch.to_worker_path(self.route.worker, is_windows=self.route.is_windows),
            "--host", host,
            "--port", str(port),
            "--token", token,
        ]
        self._process = subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

        try:
            conn, _ = listener.accept()
        except socket.timeout as exc:
            raise WorkerStartError(
                f"worker did not connect back within {self.start_timeout_s:.0f}s: "
                f"{self._drain_process()}"
            ) from exc
        finally:
            listener.close()
            self._listener = None

        conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        conn.settimeout(self.call_timeout_s)
        self._sock = conn

        op, body, _ = _proto.recv_message(conn)
        if op != _proto.OP_HELLO:
            self.close()
            raise WorkerStartError(f"first message was {op!r}, expected {_proto.OP_HELLO!r}")
        if body.get("token") != token:
            # Somebody else reached the listener first. Refuse rather than
            # speak to it: the whole point of the token.
            self.close()
            raise WorkerStartError("connect-back presented the wrong token; refusing the peer")
        self.hello = body
        return body

    def _drain_process(self) -> str:
        process = self._process
        if process is None:
            return "worker was never started"
        if process.poll() is None:
            process.kill()
        try:
            out, err = process.communicate(timeout=5)
        except subprocess.TimeoutExpired:  # pragma: no cover - kill just ran
            return f"worker pid {process.pid} did not exit"
        text = (err or b"").decode("utf-8", "replace").strip()
        if not text:
            text = (out or b"").decode("utf-8", "replace").strip()
        return f"exit={process.returncode} output={text[:2000]!r}" if text else f"exit={process.returncode}"

    def close(self, timeout_s: float = DEFAULT_CLOSE_TIMEOUT_S) -> None:
        """Ask the worker to exit; kill it if it will not. Safe to call twice."""
        sock, self._sock = self._sock, None
        if sock is not None:
            try:
                sock.settimeout(min(timeout_s, 5.0))
                _proto.send_message(sock, _proto.OP_CLOSE, {})
                _proto.recv_message(sock)
            except (OSError, ProtocolError, WorkerError):
                pass
            finally:
                sock.close()
        if self._listener is not None:
            self._listener.close()
            self._listener = None
        process, self._process = self._process, None
        if process is not None and process.poll() is None:
            try:
                process.wait(timeout=timeout_s)
            except subprocess.TimeoutExpired:
                process.kill()
                process.wait(timeout=5)
        if process is not None:
            for stream in (process.stdout, process.stderr):
                if stream is not None:
                    stream.close()

    def __enter__(self) -> ExternalWorker:
        self.start()
        return self

    def __exit__(self, *exc: object) -> None:
        self.close()

    # -- calls --------------------------------------------------------------

    def _call(
        self,
        op: str,
        body: dict[str, Any] | None = None,
        tensors: dict[str, np.ndarray] | None = None,
    ) -> tuple[dict[str, Any], dict[str, np.ndarray]]:
        sock = self._sock
        if sock is None:
            raise WorkerStartError("worker is not running; call start() first")
        try:
            _proto.send_message(sock, op, body, tensors)
            reply_op, reply_body, reply_tensors = _proto.recv_message(sock)
        except (OSError, ProtocolError) as exc:
            # The socket is the only liveness signal we have; a dead worker
            # shows up as a broken pipe or a short read, and the useful part of
            # the diagnosis is on its stderr.
            detail = self._drain_process()
            self._sock = None
            raise WorkerStartError(f"worker connection failed ({exc}); {detail}") from exc
        if reply_op != _proto.OP_OK:
            raise ProtocolError(f"unexpected reply {reply_op!r} to {op!r}")
        return reply_body, reply_tensors

    def load(
        self,
        graph_path: str | Path,
        *,
        example_inputs: dict[str, np.ndarray] | None = None,
        profile: bool = True,
        profile_prefix: str | Path | None = None,
        ep_dir: str | Path | None = None,
        intra_op_num_threads: int | None = None,
    ) -> LoadReport:
        """Open a graph and measure where its nodes actually ran.

        ``example_inputs`` are not optional in practice: without a run ORT
        writes no node assignments, so the report comes back with
        ``fraction_on_target is None`` and the planner treats the placement as
        unverified rather than assuming it is fine.
        """
        body: dict[str, Any] = {
            "graph_path": _launch.to_worker_path(graph_path, is_windows=self.route.is_windows),
            "ep": self.route.ep,
            "profile": bool(profile),
        }
        if profile_prefix is not None:
            body["profile_prefix"] = _launch.to_worker_path(
                profile_prefix, is_windows=self.route.is_windows
            )
        if ep_dir is not None:
            body["ep_dir"] = str(ep_dir)
        if intra_op_num_threads is not None:
            body["intra_op_num_threads"] = int(intra_op_num_threads)

        reply, _ = self._call(_proto.OP_LOAD, body, example_inputs or {})
        self.report = LoadReport.from_body(reply, ep=self.route.ep, graph_path=str(graph_path))
        return self.report

    def run(self, inputs: dict[str, np.ndarray]) -> tuple[dict[str, np.ndarray], RunTiming]:
        """One forward, with the transport cost separated from the compute."""
        started = time.perf_counter()
        reply, outputs = self._call(_proto.OP_RUN, {}, inputs)
        round_trip = time.perf_counter() - started
        return outputs, RunTiming(worker_s=float(reply.get("run_s", 0.0)), round_trip_s=round_trip)

    def stats(self) -> dict[str, Any]:
        """Counters and the worker's own resident set, for the ledger."""
        reply, _ = self._call(_proto.OP_STATS, {})
        return reply


def start_worker(route_name: str, **kwargs: Any) -> ExternalWorker:
    """Resolve a route by name and start it."""
    worker = ExternalWorker(_launch.resolve(route_name), **kwargs)
    worker.start()
    return worker