litellm/.github/scripts/verify_linux_native_wheel.py

from __future__ import annotations

import importlib.util
import os
import re
import subprocess
import sys
import zipfile
from collections.abc import Callable, Mapping, Sequence
from itertools import product
from pathlib import Path, PurePosixPath
from types import MappingProxyType, ModuleType
from typing import Final, Protocol

EXPECTED_PYTHON_TAG: Final = "cp310"
EXPECTED_ABI_TAG: Final = "abi3"
EXPECTED_PLATFORM_TAG: Final = "linux_x86_64"


class CommandRunner(Protocol):
    def __call__(
        self,
        command: tuple[str, ...],
        *,
        check: bool,
        capture_output: bool,
        text: bool,
    ) -> subprocess.CompletedProcess[str]: ...


def _run_command(
    command: tuple[str, ...],
    *,
    check: bool,
    capture_output: bool,
    text: bool,
) -> subprocess.CompletedProcess[str]:
    return subprocess.run(command, check=check, capture_output=capture_output, text=text)


def _dist_info_directory(member: zipfile.ZipInfo) -> str | None:
    parts: Final = PurePosixPath(member.filename).parts
    if not parts or not parts[0].endswith(".dist-info"):
        return None
    return parts[0]


def _wheel_metadata_tags(archive: zipfile.ZipFile, members: tuple[zipfile.ZipInfo, ...]) -> tuple[str, ...]:
    if len(members) != 1:
        return ()
    lines: Final = archive.read(members[0]).splitlines()
    return tuple(line.removeprefix(b"Tag:").strip().decode("ascii") for line in lines if line.startswith(b"Tag:"))


def _load_native_module(native_path: Path) -> ModuleType | None:
    module_spec: Final = importlib.util.spec_from_file_location("litellm.rust_bridge._native", native_path)
    if module_spec is None or module_spec.loader is None:
        return None
    try:
        native_module: Final = importlib.util.module_from_spec(module_spec)
        module_spec.loader.exec_module(native_module)
    except Exception as error:  # noqa: BLE001  # native module initialization can raise arbitrary exceptions
        sys.stderr.write(f"native module load failed: {error}\n")
        return None
    return native_module


def main(
    argv: Sequence[str] | None = None,
    environment: Mapping[str, str] | None = None,
    load_native_module: Callable[[Path], ModuleType | None] = _load_native_module,
    run_command: CommandRunner = _run_command,
) -> int:
    arguments: Final = tuple(sys.argv if argv is None else argv)
    resolved_environment: Final = os.environ if environment is None else environment
    gateway: Final = len(arguments) == 3 and arguments[2] == "--gateway"
    if len(arguments) != 2 and not gateway:
        sys.stderr.write(f"usage: {Path(arguments[0]).name} WHEEL [--gateway]\n")
        return 2

    expected_distribution: Final = "litellm_gateway" if gateway else "litellm"
    expected_python_tag: Final = "py3" if gateway else EXPECTED_PYTHON_TAG
    expected_abi_tag: Final = "none" if gateway else EXPECTED_ABI_TAG
    native_label: Final = "Gateway binary" if gateway else "Native extension"
    wheel: Final = Path(arguments[1])
    wheel_tags: Final = wheel.stem.rsplit("-", maxsplit=3)
    if len(wheel_tags) != 4:
        sys.stderr.write(f"cannot parse wheel tags from {wheel.name}\n")
        return 1

    wheel_identity: Final = wheel_tags[0].split("-")
    if len(wheel_identity) != 2 or wheel_identity[0] != expected_distribution or not wheel_identity[1]:
        sys.stderr.write(f"unexpected wheel identity: {wheel_tags[0]}\n")
        return 1

    expected_dist_info_directory: Final = f"{wheel_tags[0]}.dist-info"
    expected_dist_info_directories: Final = frozenset((expected_dist_info_directory,))
    python_tag: Final = wheel_tags[1]
    abi_tag: Final = wheel_tags[2]
    platform_tag: Final = wheel_tags[3]
    expanded_filename_tags: Final = frozenset(
        "-".join(tag) for tag in product(python_tag.split("."), abi_tag.split("."), platform_tag.split("."))
    )

    with zipfile.ZipFile(wheel) as archive:
        wheel_members: Final = archive.infolist()
        dist_info_directories: Final = frozenset(
            directory for member in wheel_members if (directory := _dist_info_directory(member)) is not None
        )
        required_dist_info_files: Final = ("METADATA", "RECORD", "WHEEL")
        dist_info_file_counts: Final = MappingProxyType(
            {
                filename: sum(
                    member.filename == f"{expected_dist_info_directory}/{filename}" for member in wheel_members
                )
                for filename in required_dist_info_files
            }
        )
        wheel_metadata_members: Final = tuple(
            member for member in wheel_members if member.filename == f"{expected_dist_info_directory}/WHEEL"
        )
        wheel_metadata_tags: Final = _wheel_metadata_tags(archive, wheel_metadata_members)
        native_members: Final = tuple(
            member
            for member in wheel_members
            if (
                member.filename == f"{wheel_tags[0]}.data/scripts/litellm-gateway"
                if gateway
                else member.filename.startswith("litellm/rust_bridge/_native.") and member.filename.endswith(".so")
            )
        )
        if len(native_members) != 1:
            sys.stderr.write(f"expected one {native_label.lower()}, found {len(native_members)}\n")
            return 1

        unexpected_members: Final = tuple(
            member.filename
            for member in wheel_members
            if member.filename.endswith((".pdb", ".dwp", ".rlib", ".rmeta", "Cargo.toml", "Cargo.lock"))
            or any(part.endswith(".dSYM") for part in PurePosixPath(member.filename).parts)
        )
        native_member: Final = native_members[0]
        uncompressed_wheel_size: Final = sum(member.file_size for member in wheel_members)
        native_path: Final = wheel.parent / "native" / Path(native_member.filename).name
        native_path.parent.mkdir(parents=True, exist_ok=True)
        native_bytes: Final = archive.read(native_member)
        native_path.write_bytes(native_bytes)
        duplicated_vocabularies: Final = tuple(
            member.filename
            for member in wheel_members
            if member.filename.startswith("litellm/litellm_core_utils/tokenizers/")
            and re.fullmatch(r"[0-9a-f]{40}", PurePosixPath(member.filename).name)
            and member.file_size > 0
            and archive.read(member) in native_bytes
        )

    wheel_metadata_tags_match: Final = (
        len(wheel_metadata_tags) == len(expanded_filename_tags)
        and frozenset(wheel_metadata_tags) == expanded_filename_tags
    )
    commit_sha: Final = resolved_environment.get(
        "RELEASE_WHEEL_COMMIT_SHA", resolved_environment.get("GITHUB_SHA", "unknown")
    )
    rustc_version: Final = run_command(
        ("rustc", "--version"),
        check=True,
        capture_output=True,
        text=True,
    ).stdout.strip()
    project_directory: Final = Path(__file__).parents[2] / (
        "litellm-rust/crates/gateway" if gateway else "."
    )
    pyproject: Final = (project_directory / "pyproject.toml").read_text()
    maturin_match: Final = re.search(r'"maturin==([^";]+)', pyproject)
    if maturin_match is None:
        sys.stderr.write("build-system does not pin an exact Maturin version\n")
        return 1

    maturin_version: Final = maturin_match.group(1)
    native_percentage: Final = native_member.file_size / uncompressed_wheel_size * 100
    size_report: Final = "\n".join(
        (
            "## Gateway wheel build report" if gateway else "## Native wheel build report",
            "",
            "| Build | Value |",
            "| --- | --- |",
            f"| Commit | `{commit_sha}` |",
            f"| Platform | `{platform_tag}` |",
            f"| Python ABI | `{python_tag}-{abi_tag}` |",
            f"| Rust compiler | `{rustc_version}` |",
            f"| Maturin | `{maturin_version}` |",
            "| Cargo profile | `release` |",
            "",
            "| Artifact | Size |",
            "| --- | ---: |",
            f"| Compressed wheel | {wheel.stat().st_size / 1_000_000:.2f} MB |",
            f"| Uncompressed wheel | {uncompressed_wheel_size / 1_000_000:.2f} MB |",
            f"| {native_label} | {native_member.file_size / 1_000_000:.2f} MB |",
            f"| Native share | {native_percentage:.2f}% |",
            "",
        )
    )
    summary_path: Final = resolved_environment.get("GITHUB_STEP_SUMMARY")
    if summary_path is None:
        sys.stdout.write(size_report)

    sections: Final = run_command(
        ("readelf", "--sections", "--wide", str(native_path)),
        check=True,
        capture_output=True,
        text=True,
    ).stdout
    debug_sections: Final = tuple(section for section in (".debug_", ".zdebug_") if section in sections)
    debug_sections_absent: Final = not debug_sections
    static_symbol_table_absent: Final = ".symtab" not in sections

    dynamic_symbols: Final = (
        ""
        if gateway
        else run_command(
            ("readelf", "--dyn-syms", "--wide", str(native_path)),
            check=True,
            capture_output=True,
            text=True,
        ).stdout
    )
    extension_entry_point_present: Final = "PyInit__native" in dynamic_symbols
    native_module: Final = None if gateway else load_native_module(native_path)
    native_module_loads: Final = native_module is not None
    panic_test_hook_absent: Final = native_module is not None and not hasattr(native_module, "_panic_for_test")
    native_size_limit: Final = 48_000_000
    native_size_within_limit: Final = native_member.file_size <= native_size_limit
    validations: Final = (
        (f"Python tag is {expected_python_tag}", python_tag == expected_python_tag),
        (f"ABI tag is {expected_abi_tag}", abi_tag == expected_abi_tag),
        (f"Platform tag is {EXPECTED_PLATFORM_TAG}", platform_tag == EXPECTED_PLATFORM_TAG),
        ("Wheel dist-info directory matches the filename", dist_info_directories == expected_dist_info_directories),
        (
            "Required dist-info files are present exactly once",
            all(count == 1 for count in dist_info_file_counts.values()),
        ),
        ("Wheel metadata tags match the filename", wheel_metadata_tags_match),
        ("Debug sections are absent", debug_sections_absent),
        ("Static symbol table is absent", static_symbol_table_absent),
        (f"{native_label} does not exceed {native_size_limit / 1_000_000:.0f} MB", native_size_within_limit),
        ("Tokenizer vocabularies are not duplicated in the native extension", not duplicated_vocabularies),
        ("Wheel contents are valid", not unexpected_members),
    ) + (
        ()
        if gateway
        else (
            ("Python extension entry point is present", extension_entry_point_present),
            ("Native module loads", native_module_loads),
            ("Production module omits the panic test hook", panic_test_hook_absent),
        )
    )

    verified_report: Final = size_report + "\n".join(
        ("", "| Validation | Expected | Result |", "| --- | --- | :---: |")
        + tuple(f"| {label} | Yes | {'O' if passed else 'X'} |" for label, passed in validations)
        + ("",)
    )
    if summary_path is not None:
        with Path(summary_path).open("a") as summary:
            summary.write(verified_report + "\n")

    invalid_dist_info_files: Final = any(count != 1 for count in dist_info_file_counts.values())
    validation_errors: Final = tuple(
        message
        for failed, message in (
            (bool(debug_sections), f"{native_member.filename} contains debug sections: {', '.join(debug_sections)}"),
            (not static_symbol_table_absent, f"{native_member.filename} contains a static symbol table"),
            (not gateway and not extension_entry_point_present, "native extension does not export PyInit__native"),
            (
                python_tag != expected_python_tag,
                f"unexpected Python tag: expected {expected_python_tag}, found {python_tag}",
            ),
            (abi_tag != expected_abi_tag, f"unexpected ABI tag: expected {expected_abi_tag}, found {abi_tag}"),
            (
                platform_tag != EXPECTED_PLATFORM_TAG,
                f"unexpected platform tag: expected {EXPECTED_PLATFORM_TAG}, found {platform_tag}",
            ),
            (
                dist_info_directories != expected_dist_info_directories,
                f"unexpected dist-info directories: expected {expected_dist_info_directory}, "
                f"found {', '.join(sorted(dist_info_directories))}",
            ),
            (invalid_dist_info_files, f"required dist-info file counts are invalid: {dist_info_file_counts}"),
            (
                not invalid_dist_info_files and not wheel_metadata_tags_match,
                f"WHEEL tags do not match filename: expected {', '.join(sorted(expanded_filename_tags))}, "
                f"found {', '.join(sorted(wheel_metadata_tags))}",
            ),
            (
                native_module is not None and not panic_test_hook_absent,
                "production native module exposes _panic_for_test",
            ),
            (
                not native_size_within_limit,
                f"{native_label.lower()} exceeds {native_size_limit / 1_000_000:.0f} MB: "
                f"{native_member.file_size / 1_000_000:.2f} MB",
            ),
            (bool(unexpected_members), f"wheel contains unexpected build artifacts: {', '.join(unexpected_members)}"),
        )
        if failed
    )
    sys.stderr.write("".join(f"{message}\n" for message in validation_errors))

    return 0 if all(passed for _, passed in validations) else 1


if __name__ == "__main__":
    sys.exit(main())