InfiniLM-ModelHub/tools/build_backend_plugins.py at master · InfiniTensor/InfiniLM-ModelHub · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
#!/usr/bin/env python3
from __future__ import annotations

import argparse
import os
import subprocess
from pathlib import Path


def _find_one(pattern: str, root: Path) -> Path:
    matches = sorted(root.glob(pattern))
    if not matches:
        raise FileNotFoundError(f"no file matched {pattern!r} under {root}")
    return matches[0]


def main() -> None:
    parser = argparse.ArgumentParser(description="Build InfiniLM-ModelHub C++ backend plugins.")
    parser.add_argument(
        "--infinilm-root",
        default=os.environ.get("INFINILM_REPO"),
        help="Path to an InfiniLM source checkout. Can also be set with INFINILM_REPO.",
    )
    parser.add_argument(
        "--infini-root",
        default=os.environ.get("INFINI_ROOT"),
        help="Path to the InfiniCore installation prefix. Can also be set with INFINI_ROOT.",
    )
    parser.add_argument(
        "--source",
        default="cpp_backends/modelhub_backend_adapters.cpp",
        help="C++ source file for the backend plugin, relative to this repo unless absolute.",
    )
    parser.add_argument(
        "--output-name",
        default="libinfinilm_modelhub_backends.so",
        help="Output shared library name.",
    )
    parser.add_argument("--output-dir", default="build/backend_plugins")
    parser.add_argument("--cxx", default=os.environ.get("CXX", "g++"))
    args = parser.parse_args()
    if not args.infinilm_root:
        parser.error("--infinilm-root or INFINILM_REPO is required")
    if not args.infini_root:
        parser.error("--infini-root or INFINI_ROOT is required")

    repo_root = Path(__file__).resolve().parents[1]
    infinilm_root = Path(args.infinilm_root).expanduser().resolve()
    infini_root = Path(args.infini_root).expanduser().resolve()
    output_dir = (repo_root / args.output_dir).resolve()
    output_dir.mkdir(parents=True, exist_ok=True)

    infinilm_lib_dir = infinilm_root / "python" / "infinilm" / "lib"
    infinilm_extension = _find_one("_infinilm*.so", infinilm_lib_dir)
    output_so = output_dir / args.output_name
    source = Path(args.source).expanduser()
    if not source.is_absolute():
        source = repo_root / source
    source = source.resolve()

    command = [
        args.cxx,
        "-std=c++17",
        "-O2",
        "-fPIC",
        "-shared",
        str(source),
        "-o",
        str(output_so),
        "-I",
        str(infinilm_root / "csrc"),
        "-I",
        str(infinilm_root / "third_party" / "json" / "single_include"),
        "-I",
        str(infinilm_root / "third_party" / "spdlog" / "include"),
        "-I",
        str(infini_root / "include"),
        str(infinilm_extension),
        "-Wl,-rpath," + str(infinilm_lib_dir),
        "-Wl,-rpath," + str(infini_root / "lib"),
        "-ldl",
    ]

    subprocess.run(command, check=True)
    print(output_so)


if __name__ == "__main__":
    main()