overlay/scripts/bootstrap_benchmark_runtime.py

#!/usr/bin/env python3
from __future__ import annotations

import importlib.util
import json


PACKAGE_MAP = {
    "mamba_ssm": "mamba-ssm",
    "transformers": "transformers",
}


def detect_missing_modules(required: list[str] | None = None) -> list[str]:
    names = required or list(PACKAGE_MAP)
    return [name for name in names if importlib.util.find_spec(name) is None]


def build_install_command(*, missing_modules: list[str]) -> list[str]:
    packages = [PACKAGE_MAP[name] for name in missing_modules if name in PACKAGE_MAP]
    if not packages:
        return []
    return ["python", "-m", "pip", "install", *packages]


def build_runtime_report(*, missing_modules: list[str]) -> dict[str, object]:
    return {
        "ready": len(missing_modules) == 0,
        "missing_modules": list(missing_modules),
        "packages": {name: PACKAGE_MAP[name] for name in missing_modules if name in PACKAGE_MAP},
        "install_command": build_install_command(missing_modules=missing_modules),
    }


def main() -> int:
    missing = detect_missing_modules()
    print(json.dumps(build_runtime_report(missing_modules=missing), indent=2, sort_keys=True))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())