litellm/tests/integration/_support/process.py

import os
import signal
import socket
import subprocess
import sys
import time
import uuid
from collections.abc import Iterator, Mapping
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from types import MappingProxyType
from typing import Final

import httpx
import psutil
from integration._support.client import Gateway


def proxy_database_environment() -> Mapping[str, str]:
    writer: Final = os.environ.get("INTEGRATION_PROXY_DATABASE_URL", "")
    reader: Final = os.environ.get("INTEGRATION_PROXY_READ_REPLICA_URL", "")
    return MappingProxyType(
        {
            **({"DATABASE_URL": writer} if writer else {}),
            **({"DATABASE_URL_READ_REPLICA": reader} if reader else {}),
        }
    )


def in_group(process: psutil.Process, group: int) -> bool:
    try:
        return os.getpgid(process.pid) == group
    except ProcessLookupError:
        return False


def group_members(group: int) -> tuple[psutil.Process, ...]:
    return tuple(process for process in psutil.process_iter() if in_group(process, group))


def signal_group(group: int, action: int) -> None:
    try:
        os.killpg(group, action)
    except ProcessLookupError:
        pass


def stop_root_process(process: subprocess.Popen[bytes]) -> bool:
    if process.poll() is not None:
        return True
    process.terminate()
    try:
        process.wait(timeout=30)
    except subprocess.TimeoutExpired:
        return False
    return True


@dataclass(frozen=True, slots=True)
class OwnedProxy:
    gateway: Gateway
    process: subprocess.Popen[bytes]
    log: Path


@contextmanager
def owned_proxy(
    gateway: Gateway,
    directory: Path,
    overrides: Mapping[str, str],
    *,
    config: Path | None = None,
    remove_environment: tuple[str, ...] = (),
    workers: int = 1,
) -> Iterator[Gateway]:
    with owned_proxy_process(
        gateway, directory, overrides, config=config, remove_environment=remove_environment, workers=workers
    ) as owned:
        yield owned.gateway


@dataclass(frozen=True, slots=True)
class LaunchedProxy:
    port: int
    process: subprocess.Popen[bytes]
    log: Path


@contextmanager
def launched_proxy(
    gateway: Gateway,
    directory: Path,
    overrides: Mapping[str, str],
    *,
    config: Path | None = None,
    remove_environment: tuple[str, ...] = (),
    workers: int = 1,
) -> Iterator[LaunchedProxy]:
    with socket.socket() as reserve:
        reserve.bind(("127.0.0.1", 0))
        port: Final = reserve.getsockname()[1]
    root: Final = Path(os.environ.get("INTEGRATION_PROXY_ROOT") or Path(__file__).resolve().parents[3])
    environment: Final = {
        **{
            name: value
            for name, value in {**os.environ, **proxy_database_environment()}.items()
            if name not in remove_environment
        },
        "LITELLM_MASTER_KEY": gateway.key,
        "LITELLM_SALT_KEY": os.environ.get("LITELLM_SALT_KEY", "sk-integration-salt"),
        "STORE_MODEL_IN_DB": "True",
        **overrides,
    }
    output: Final = Path(os.environ.get("INTEGRATION_RESULTS_DIR", str(directory)))
    output.mkdir(parents=True, exist_ok=True)
    log_path: Final = output / f"owned-proxy-{uuid.uuid4().hex}.log"
    with log_path.open("w") as log:
        process: Final = subprocess.Popen(
            [
                sys.executable,
                "-m",
                "integration._support.proxy",
                "--config",
                str(config or "tests/integration/proxy_config.yaml"),
                "--host",
                "127.0.0.1",
                "--port",
                str(port),
                "--num_workers",
                str(workers),
                "--use_prisma_db_push",
                "--enforce_prisma_migration_check",
            ],
            cwd=root,
            env=environment,
            stdout=log,
            stderr=subprocess.STDOUT,
            start_new_session=True,
        )
        try:
            yield LaunchedProxy(port, process, log_path)
        finally:
            root_stopped: Final = stop_root_process(process)
            residual: Final = group_members(process.pid)
            if residual:
                signal_group(process.pid, signal.SIGTERM)
                psutil.wait_procs(residual, timeout=5)
            remaining: Final = group_members(process.pid)
            if remaining:
                signal_group(process.pid, signal.SIGKILL)
                psutil.wait_procs(remaining, timeout=3)
            process.wait(timeout=3)
            survivors: Final = group_members(process.pid)
            assert not survivors, "Owned proxy child survived cleanup"
            assert root_stopped and not remaining, "Owned proxy required forced cleanup"


def _is_ready(client: httpx.Client) -> bool:
    try:
        return client.get("/health/readiness", timeout=2).status_code == 200
    except httpx.TransportError:
        return False


def refused_boot_log(
    gateway: Gateway,
    directory: Path,
    overrides: Mapping[str, str],
    *,
    config: Path | None = None,
) -> str:
    """Start the proxy and return its log once it exits non-zero instead of becoming ready."""
    with launched_proxy(gateway, directory, overrides, config=config) as launched:
        with httpx.Client(base_url=f"http://127.0.0.1:{launched.port}", timeout=15, trust_env=False) as client:
            deadline: Final = time.monotonic() + 70
            while launched.process.poll() is None:
                assert not _is_ready(client), (
                    f"Proxy became ready instead of refusing to boot:\n{launched.log.read_text()}"
                )
                assert time.monotonic() < deadline, "Proxy neither exited nor became ready within the deadline"
                time.sleep(0.1)
        assert launched.process.returncode != 0, (
            f"Proxy exited 0 instead of refusing to boot:\n{launched.log.read_text()}"
        )
        return launched.log.read_text()


@contextmanager
def owned_proxy_process(
    gateway: Gateway,
    directory: Path,
    overrides: Mapping[str, str],
    *,
    config: Path | None = None,
    remove_environment: tuple[str, ...] = (),
    workers: int = 1,
) -> Iterator[OwnedProxy]:
    with launched_proxy(
        gateway, directory, overrides, config=config, remove_environment=remove_environment, workers=workers
    ) as launched:
        with httpx.Client(base_url=f"http://127.0.0.1:{launched.port}", timeout=15, trust_env=False) as client:
            deadline: Final = time.monotonic() + 70
            while True:
                assert launched.process.poll() is None, "Owned proxy exited before readiness"
                try:
                    if client.get("/health/readiness", timeout=2).status_code == 200:
                        break
                except httpx.TransportError:
                    pass
                assert time.monotonic() < deadline, "Owned proxy readiness deadline exceeded"
                time.sleep(0.1)
            yield OwnedProxy(Gateway(client, gateway.key, gateway.upstream_url), launched.process, launched.log)