forked from ccf-ai-infra/Intro-ops
99 lines
3.2 KiB
Python
99 lines
3.2 KiB
Python
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
try:
|
|
import yaml
|
|
except ImportError as exc: # pragma: no cover
|
|
raise SystemExit("PyYAML is required to generate operator artifacts") from exc
|
|
|
|
|
|
def load_operator_manifests(ops_root: Path) -> list[dict[str, Any]]:
|
|
manifests: list[dict[str, Any]] = []
|
|
for path in sorted(ops_root.glob("*/operator.yaml")):
|
|
data = yaml.safe_load(path.read_text()) or {}
|
|
data["_path"] = str(path)
|
|
manifests.append(data)
|
|
return manifests
|
|
|
|
|
|
def cmake_bool(name: str) -> str:
|
|
return f"${{{name}}}"
|
|
|
|
|
|
def generate_operators_cmake(manifests: list[dict[str, Any]], ops_root: Path) -> str:
|
|
lines = [
|
|
"# Generated by tools/generate_operator_artifacts.py",
|
|
"set(CAMP_OPERATOR_SOURCES",
|
|
" ${CMAKE_CURRENT_SOURCE_DIR}/common/status.cc",
|
|
]
|
|
for manifest in manifests:
|
|
nvidia = manifest.get("backends", {}).get("nvidia", {})
|
|
if nvidia.get("status") != "runnable":
|
|
continue
|
|
for source in nvidia.get("sources", []):
|
|
relative = Path(source)
|
|
if relative.parts and relative.parts[0] == "ops":
|
|
relative = Path(*relative.parts[1:])
|
|
lines.append(f" $<$<BOOL:{cmake_bool('CAMP_ENABLE_NVIDIA')}>:${{CMAKE_CURRENT_SOURCE_DIR}}/{relative.as_posix()}>")
|
|
lines.append(")")
|
|
lines.append("")
|
|
return "\n".join(lines)
|
|
|
|
|
|
def python_literal(value: Any) -> str:
|
|
return repr(value)
|
|
|
|
|
|
def generate_python_registry(manifests: list[dict[str, Any]]) -> str:
|
|
serializable = []
|
|
for manifest in manifests:
|
|
item = {key: value for key, value in manifest.items() if not key.startswith("_")}
|
|
serializable.append(item)
|
|
return "\n".join(
|
|
[
|
|
"# Generated by tools/generate_operator_artifacts.py",
|
|
"from __future__ import annotations",
|
|
"",
|
|
"OPERATORS = " + python_literal(serializable),
|
|
"",
|
|
"def get_operator(name: str):",
|
|
" for op in OPERATORS:",
|
|
" if op['name'] == name:",
|
|
" return op",
|
|
" raise KeyError(name)",
|
|
"",
|
|
]
|
|
)
|
|
|
|
|
|
def generate_test_manifest(manifests: list[dict[str, Any]]) -> str:
|
|
payload = [{key: value for key, value in manifest.items() if not key.startswith("_")} for manifest in manifests]
|
|
return json.dumps(payload, indent=2, sort_keys=True) + "\n"
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--ops-root", type=Path, required=True)
|
|
parser.add_argument("--out-dir", type=Path, required=True)
|
|
args = parser.parse_args()
|
|
|
|
manifests = load_operator_manifests(args.ops_root)
|
|
if not manifests:
|
|
raise SystemExit(f"no operator manifests found under {args.ops_root}")
|
|
|
|
args.out_dir.mkdir(parents=True, exist_ok=True)
|
|
(args.out_dir / "operators.cmake").write_text(generate_operators_cmake(manifests, args.ops_root))
|
|
(args.out_dir / "operator_registry.py").write_text(generate_python_registry(manifests))
|
|
(args.out_dir / "operator_test_manifest.json").write_text(generate_test_manifest(manifests))
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|
|
|