mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-30 06:43:27 +00:00
75 lines
2.5 KiB
Python
75 lines
2.5 KiB
Python
from __future__ import annotations
|
|
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import zipfile
|
|
from pathlib import Path
|
|
from typing import Final
|
|
|
|
COST_MAP: Final = Path(__file__).resolve().parents[2] / "model_prices_and_context_window.json"
|
|
|
|
CHILD_SCRIPT: Final = """
|
|
from importlib.util import module_from_spec, spec_from_file_location
|
|
from pathlib import Path
|
|
import sys
|
|
|
|
native_path = Path(sys.argv[1])
|
|
spec = spec_from_file_location("litellm.rust_bridge._native", native_path)
|
|
if spec is None or spec.loader is None:
|
|
raise RuntimeError("cannot create native extension import specification")
|
|
module = module_from_spec(spec)
|
|
spec.loader.exec_module(module)
|
|
|
|
before = module.gil_stats()
|
|
if not isinstance(before.get("releases"), int):
|
|
raise AssertionError(f"unexpected gil_stats result: {before!r}")
|
|
|
|
try:
|
|
module._panic_for_test()
|
|
except BaseException as error:
|
|
if type(error).__name__ != "PanicException":
|
|
raise AssertionError(f"expected PanicException, got {type(error).__name__}") from error
|
|
else:
|
|
raise AssertionError("Rust panic returned without raising")
|
|
|
|
after = module.gil_stats()
|
|
if not isinstance(after.get("releases"), int):
|
|
raise AssertionError(f"native module unusable after panic: {after!r}")
|
|
|
|
if module.bundled_model_cost_map() != Path(sys.argv[2]).read_bytes():
|
|
raise AssertionError("the wheel embeds a different model_prices_and_context_window.json than this checkout")
|
|
"""
|
|
|
|
|
|
def main() -> int:
|
|
if len(sys.argv) != 2:
|
|
sys.stderr.write(f"usage: {Path(sys.argv[0]).name} WHEEL\n")
|
|
return 2
|
|
|
|
wheel: Final = Path(sys.argv[1])
|
|
with tempfile.TemporaryDirectory() as temporary_directory, zipfile.ZipFile(wheel) as archive:
|
|
native_members: Final = tuple(
|
|
member
|
|
for member in archive.infolist()
|
|
if member.filename.startswith("litellm/rust_bridge/_native.") and member.filename.endswith(".so")
|
|
)
|
|
if len(native_members) != 1:
|
|
sys.stderr.write(f"expected one native extension, found {len(native_members)}\n")
|
|
return 1
|
|
|
|
native_path: Final = Path(temporary_directory) / Path(native_members[0].filename).name
|
|
native_path.write_bytes(archive.read(native_members[0]))
|
|
result: Final = subprocess.run(
|
|
(sys.executable, "-c", CHILD_SCRIPT, str(native_path), str(COST_MAP)), check=False
|
|
)
|
|
|
|
if result.returncode != 0:
|
|
sys.stderr.write(f"native wheel smoke test exited with status {result.returncode}\n")
|
|
return 1
|
|
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|